python - 使用 np.diag 分配移动的对角线值

Question

我正在尝试为矩阵的多个对角线赋值。例如，我有这个矩阵:

>>> u = np.zeros(25).reshape(5, 5)
>>> u
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]])

我想为上面的第 $k$ 条对角线赋值。例如，如果 $k=1$，我想要主对角线上方的对角线。我试图通过使用像这样的 np.diag 来实现这一点 np.diag(u, k=1) = 1 我想得到以下结果:

>>> u
array([[0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.],
       [0., 0., 0., 0., 0.]])

问题是这会抛出一个 SyntaxError: can't assign to function call 因为这又是 Python。但是 np.diag 返回对原始矩阵的引用，如您所见:

>>> np.may_share_memory(np.diag(u, k=1), u)
True

我该怎么做？提前谢谢你。

Answer 1

你可以使用

u[np.eye(len(u), k=1, dtype='bool')] = 1
print(u)

输出:

[[0. 1. 0. 0. 0.]
 [0. 0. 1. 0. 0.]
 [0. 0. 0. 1. 0.]
 [0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0.]]