Python:单独取矩阵每一行的外积，取和然后返回一个和向量

Question

假设我有一个 N 乘 M 的矩阵 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)