python - 如何只添加到 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]])

由于我可能还要处理非常大的矩阵数组， 按照此 thread 的公认解决方案中的建议，通过赋值操作执行此对角线加法的最有效方法是什么？ ？

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()来得到扁平化 View 然后赋值-

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