Python：取每行矩阵的外积，取总和，然后返回和的向量

Question

假设我有一个维数为N的矩阵A.

我希望返回N维向量V，其中第n个元素是A的第n行中条目的所有成对乘积的双和。

在循环中，我想我能做到：

V = np.zeros(A.shape[0])
for n in range(A.shape[0]):
    for i in range(A.shape[1]):
        for j in range(A.shape[1]):
            V[n] += A[n,i] * A[n,j]

我想对此进行矢量化，我想我可以做到：

V_temp = np.einsum('ij,ik->ijk', A, A)
V = np.einsum('ijk->i', A)

但是我不认为这是非常有效的内存方式，因为当我需要的是总和时，中间步骤V_temp不必要地存储整个外部产品。有更好的方法吗？

由于

Answer 1

您可以使用

V=np.einsum("ni,nj->n",A,A)

Answer 2

你实际上正在计算

A.sum(-1)**2

换句话说，外部产品的总和只是因子总和的乘积。

演示：

A = np.random.random((1000,1000))
np.allclose(np.einsum('ij,ik->i', A, A), A.sum(-1)**2)
# True
t = timeit.timeit('np.einsum("ij,ik->i",A,A)', globals=dict(A=A,np=np), number=10)*100; f"{t:8.4f} ms"
# '948.4210 ms'
t = timeit.timeit('A.sum(-1)**2', globals=dict(A=A,np=np), number=10)*100; f"{t:8.4f} ms"
# '  0.7396 ms'

Answer 3

也许你可以使用

np.einsum('ij,ik->i', A, A)

或等效的

np.einsum(A, [0,1], A, [0,2], [0])

在2015年的Macbook上，我得到了

In [35]: A = np.random.rand(100,100)

In [37]: %timeit for_loops(A)
640 ms ± 24.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [38]: %timeit np.einsum('ij,ik->i', A, A)
658 µs ± 7.25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [39]: %timeit np.einsum(A, [0,1], A, [0,2], [0])
672 µs ± 19.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)