python - 在 Numpy 中制作特殊的对角矩阵

Question

我正在尝试制作一个如下所示的 numpy 数组:

[a b c       ]
[  a b c     ]
[    a b c   ]
[      a b c ] 

所以这涉及到更新主对角线和它上面的两条对角线。

执行此操作的有效方法是什么？

Answer 1

您可以使用 np.indices 获取数组的索引，然后将值分配到您想要的位置。

a = np.zeros((5,10))
i,j = np.indices(a.shape)

i,j 分别是行索引和列索引。

a[i==j] = 1.
a[i==j-1] = 2.
a[i==j-2] = 3.

将导致:

array([[ 1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.]])