我在 Python 上有一个数组，如下所示：
array([[ 0.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 0.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 0.86064797]])
我想在数组的对角线上添加一个标量值 20，使得输出为：
array([[ 20.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 20.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 20.86064797]])
由于我可能还要处理非常大的矩阵数组， 正如thread 的公认解决方案中所建议的那样，通过赋值操作进行对角相加的最有效方法是什么？
【问题讨论】：
一种方法是使用适当的步长分配扁平切片 -
In [233]: a
Out[233]:
array([[ 0.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 0.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 0.86064797]])
In [234]: a.flat[::a.shape[1]+1] += 20
In [235]: a
Out[235]:
array([[ 20.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 20.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 20.86064797]])
我们也可以使用ndarray.ravel()获取扁平化视图然后赋值-
a.ravel()[::a.shape[1]+1] += 20
另一种方法是使用np.einsum，它可以让我们看到对角线元素 -
In [269]: a
Out[269]:
array([[ 0.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 0.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 0.86064797]])
In [270]: d = np.einsum('ii->i', a)
In [271]: d += 20
In [272]: a
Out[272]:
array([[ 20.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 20.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 20.86064797]])
In [285]: a = np.random.rand(10000,10000)
# @Willem Van Onsem's soln
In [286]: %timeit np.fill_diagonal(a, a.diagonal() + 20)
10000 loops, best of 3: 159 µs per loop
In [287]: %timeit a.flat[::a.shape[1]+1] += 20
10000 loops, best of 3: 179 µs per loop
In [288]: %timeit a.ravel()[::a.shape[1]+1] += 20
100000 loops, best of 3: 18.2 µs per loop
In [289]: %%timeit
...: d = np.einsum('ii->i', a)
...: d += 20
100000 loops, best of 3: 18.5 µs per loop
【讨论】：
鉴于a是我们要更新的数组，我们可以使用.diagonal()和np.fill_diagonal：
np.fill_diagonal(a, a.diagonal() + 20)
因此，我们首先使用a.diagonal() 获取a 的对角线，然后将20 添加到对角线的每个元素。我们使用np.fill_diagonal(..)来设置对角线的元素。
【讨论】：
您还可以使用np.eye 创建一个对角矩阵，并将其与合适的常数c 相乘。
np.eye(n, dtype=int) * c # n x n matrix with diagonal part being c
然后，将其添加到您的矩阵 a
a = np.array([[ 0.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 0.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 0.86064797]])
a += np.eye(3) * 20
或者也许用np.eye(3, dtype=bool) 制作一个蒙版，这样你就可以用它来选择对角线部分并像这样添加它
a[np.eye(3, dtype=bool)] += 20
感谢 Divakar 对使用面罩的建议以及 Willem Van Onsem 的建议以明确答案。
编辑：Willem 和 Divakar 的方法在数据量很大时都比这快很多。
【讨论】：
a += ，以便清楚您更新a。
a[np.eye(n, dtype=bool)] += 20。用于掩蔽的布尔数组对于大型矩阵形状的内存将不那么重要。
注意到您的 nD 数组是正方形（或一般情况下的长方体），您可以提取对角元素的索引np.diag_indices_from，通过索引到数组中获取对角元素的原始值，然后添加您想要的 常数;然后通过使用我们从np.diag_indices 获得的索引对数组进行索引来更新原始对角线值，例如：
In [28]: arr
Out[28]:
array([[ 0.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 0.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 0.86064797]])
In [29]: new_diag_vals = arr[np.diag_indices_from(arr)] + 20
In [30]: arr[np.diag_indices(arr.shape[0])] = new_diag_vals
In [31]: arr
Out[31]:
array([[ 20.57733218, 0.09794384, 0.44497735],
[ 0.87061284, 20.10253493, 0.56643557],
[ 0.76358739, 0.44902046, 20.86064797]])
如果你想在一行中完成，那么：
In [36]: arr[np.diag_indices(arr.shape[0])] = arr[np.diag_indices_from(arr)] + 20
P.S.请注意，这会修改原始数组。
【讨论】：