python - numpy:用不同的值填充偏移对角线

Question

我需要制作一个 n*n 矩阵 m 其元素遵循 m(i,i+1)=sqrt(i)否则为 0。例如，对于 n=5，我们应该有

[0 a 0 0 0]
[0 0 b 0 0]
[0 0 0 c 0]
[0 0 0 0 d]
[0 0 0 0 0]

其中 {a,b,c,d}=sqrt({1,2,3,4})。 Here是常量三对角矩阵的解，但我的情况比这复杂一点。我知道我可以用循环或列表理解来做到这一点，但还有其他方法吗？ n 可能很大。

例如(列表理解代码)

ele=lambda i,j:sqrt(i+1) if j-i==1 else 0

[[ele(i,j) for j in range(0,6)] for i in range(0,6)]

Answer 1

你可以直接使用np.diag:

>>> d = np.sqrt(1 + np.arange(4))
>>> np.diag(d, 1)
array([[ 0.        ,  1.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  1.41421356,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  1.73205081,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  2.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ]])

np.diag 的第二个参数指定了所讨论的对角线。