python - 创建一个对角矩阵

为了方便地创建双对角矩阵或三对角矩阵，numpy或scipy中是否有特定的函数？
到目前为止，我想出的最简单的解决方案是：

main_diag = [1,2,3,4,5]
off1 = [1,2,3,4]
np.diag(main_diag) + np.diag(off1, 1)

array([[1, 1, 0, 0, 0],
       [0, 2, 2, 0, 0],
       [0, 0, 3, 3, 0],
       [0, 0, 0, 4, 4],
       [0, 0, 0, 0, 5]])

有什么功能可以一步完成吗？

最佳答案

您可以使用scipy中的sparse.diags。这里有一个简单的操作解决方案。

from scipy import sparse

diags = [range(1, 6), range(1, 5)]
sparse.diags(diags, [0, 1]).toarray()

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

概括来说，

size = 5
num_diags = 2 # should be < size
diags = [range(1, size+1-i) for i in range(num_diags)]

sparse.diags(diags, range(num_diags)).toarray()

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