如何在numpy中更改矩阵对角线的值？

Question

我想将某个numpy矩阵的对角线设置为任意一维数组。

例如，如果：

mat = np.array([[1,2],[3,4]])
diag = np.array([5,6])

然后：

>>> set_diagonal(mat, diag)
>>> mat
... array([[5,2],
           [3,6]])

我可以从diag变量，diag_fill mat矩阵到0创建对角矩阵，然后将结果相加。但是，如果我要更改对角线，这似乎会使用大量不必要的资源。

如何在numpy中设置矩阵的对角线？

注意：这不是this question的副本，因为他们想将对角线设置为常数。

Answer 1

您正在寻找numpy.fill_diagonal

根据文档：

  

numpy.fill_diagonal（a，val，wrap = False）。
  填充任意维度的给定数组的主对角线。

import numpy as np
mat = np.array([[1,2],[3,4]])

np.fill_diagonal(mat, [5,6])
print(mat)
#[[5 2]
# [3 6]]

或

import numpy as np
mat = np.array([[1,2],[3,4]])
diag = np.array([5,6])
np.fill_diagonal(mat, diag)
print(mat)
#[[5 2]
# [3 6]]

Answer 2

您可以将该1d数组传递给fill_diagonal。

np.fill_diagonal(mat, [5,6])

Answer 3

使用fill_diagonal的答案更加优雅并且计算效率更高，但是您也可以访问对角线索引然后进行更改：

mat = np.array([[1,2],[3,4]])
diag = np.array([5,6])

row, col = np.diag_indices(mat.shape[0])

mat[row, col] = diag

有关更多信息，请参见here：