対角成分以外を0にする行列

Question

対角成分以外を0とするプログラムを作りたいです

def trans(a):
    n=a.shape[0]
    for i in range(n):
        for j in range(i,n):
            tmp=a[i][j]
            a[i][j]=0
            a[j][i]=0
            a[j][i]=tmp       
a=np.array([[2,5,3],[1,6,4],[9,3,8]])
print('original matrix');print(a)
trans(a)
print('transposed matrix');print(a)

この状態で実行すると
original matrix
[[2 5 3]
[1 6 4]
[9 3 8]]
transposed matrix
[[2 0 0]
[5 6 0]
[3 4 8]]　と表示されてしまい
求めている
[[2 0 0]
[0 6 0]
[0 0 8]] という形になりません。
また、出来ればfor文で作りたいです。

Answer 1

iとjが異なるときに0をセットすればよいです。

def trans(a):
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            if i != j:
                a[i][j] = 0

Answer 2

下記のような感じでしょうか。for文ではないですが…

import numpy as np
A = np.array([[2,5,3],[1,6,4],[9,3,8]])
print(A)
"""
[[2 5 3]
 [1 6 4]
 [9 3 8]]
"""
E = np.identity(3)
print(E)
"""
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
"""
AE = A*E
print(AE)
"""
[[2. 0. 0.]
 [0. 6. 0.]
 [0. 0. 8.]]
"""