python - 如何在python中使用for循环转换单位矩阵

Question

所以，假设我有一个单位矩阵 AB，我想把对角线的变成 5，把对角线的直接变成 8s

    AB=np.identity(5)
    AB
    array([[1., 0., 0., 0., 0.],
           [0., 1., 0., 0., 0.],
           [0., 0., 1., 0., 0.],
           [0., 0., 0., 1., 0.],
           [0., 0., 0., 0., 1.]])
row=AB.shape[0]
col=AB.shape[1]
new = AB
for i in range (0,row):
    for j in range (0,col): 
        if AB[i,j] !=0:
            new[i,j] = 5

       ## if j+1 == 0:
           ## new[i,j+1] = 8 
            
print(new)
[[5. 0. 0. 0. 0.]
 [0. 5. 0. 0. 0.]
 [0. 0. 5. 0. 0.]
 [0. 0. 0. 5. 0.]
 [0. 0. 0. 0. 5.]]

使用上面的代码，我可以将 1 更改为 5，但我无法弄清楚如何将主对角线两侧的第一个对角线更改为 8

标签： python

Answer 1

您可以使用以下条件：

if abs(i-j)==1: new[i,j] = 8

或者，您可以直接使用值构建数组：

import numpy as np

N = 5
AB = np.array(([5,8,*[0]*(N-2),8]*N)[:N*N]).reshape((N,N))
print(AB)

[[5 8 0 0 0]
 [8 5 8 0 0]
 [0 8 5 8 0]
 [0 0 8 5 8]
 [0 0 0 8 5]]

另一种选择是使用范围分配：

import numpy as np

N = 5
AB=np.identity(N)*5
AB[range(N-1),range(1,N)]=8
AB[range(1,N),range(N-1)]=8

print(AB)
[[5. 8. 0. 0. 0.]
 [8. 5. 8. 0. 0.]
 [0. 8. 5. 8. 0.]
 [0. 0. 8. 5. 8.]
 [0. 0. 0. 8. 5.]]