Numpy将值添加到对角线

时间:2018-08-10 14:25:17

标签: python numpy matrix multidimensional-array

我希望创建一个NxN numpy数组,其中仅填充主对角线上的所有内容。 填充的方式是主对角线(k=0gamma**0填充,k=1对角线用gamma**1填充,k=2对角线是充满gamma**2等...

gamma = 0.9
dim = 4

M = np.zeros((dim,dim))
for i in range(dim)[::-1]:   
    M += np.diagflat([gamma**(dim-i-1)]*(i+1),dim-i-1) 

print(M)

这正确地给出了

array([[ 1.   ,  0.9  ,  0.81 ,  0.729],
       [ 0.   ,  1.   ,  0.9  ,  0.81 ],
       [ 0.   ,  0.   ,  1.   ,  0.9  ],
       [ 0.   ,  0.   ,  0.   ,  1.   ]])

我想问一下是否有其他更简单或更优雅的方式来处理此问题或其他不同方法。我将经常处理多维数组,我想从不同的工具和方法中获得启发。

1 个答案:

答案 0 :(得分:4)

一种方法是使用np.triu_indices创建上三角索引,然后使用advanced indexing将值分配给这些位置:

M = np.zeros((dim,dim))

rowidx, colidx = np.triu_indices(dim)
# the diagonal offset can be calculated by subtracting the row index from column index
M[rowidx, colidx] = gamma ** (colidx - rowidx) 

M
#array([[ 1.   ,  0.9  ,  0.81 ,  0.729],
#       [ 0.   ,  1.   ,  0.9  ,  0.81 ],
#       [ 0.   ,  0.   ,  1.   ,  0.9  ],
#       [ 0.   ,  0.   ,  0.   ,  1.   ]])