如何在Python中只添加数组的对角线？

Question

我在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]])

由于我也可能正在处理非常大的矩阵数组,通过分配操作进行这种对角线加法的最有效方法是什么,如该线程的公认解决方案中所建议的那样？

Answer 1

一种方法是在具有适当步长的扁平切片上进行分配 -

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

Answer 2

鉴于这a是我们想要更新的数组,我们可以使用.diagonal()和np.fill_diagonal:

np.fill_diagonal(a, a.diagonal() + 20)

因此,我们首先获取awith 的对角线a.diagonal(),然后添加20到对角线的每个元素.我们np.fill_diagonal(..)用来设置对角线的元素.

Answer 3

您还可以创建一个对角矩阵，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 的方法都比这快得多。