python - Numpy 整个矩阵与单行相乘

Question

我有两个 numpy 矩阵，A 和 B:

A = [[0,-1],
     [-2,5],
     [2,6]]
B = [[5,3],
     [2,4]]
Dimension of A: N * K
Dimension of B: K * K

我想做的是与A的每一行和整个B进行矩阵乘法，然后与A的每一行相乘 再次。然后总结一下。

这是:

N = A.shape[0]
sum = 0
for i in range(N):
    s = A[i] @ B @ A[i]
    print(s)
    sum += s
return sum

此打印

4
70
224

而 4 + 70 + 224 = 298 就是我想要的。 有没有任何优雅或Python式的方法可以在没有for循环的情况下实现这一点？ 谢谢

Answer 1

我们可以使用np.einsum -

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

并使用optimize标志进一步优化它 -

np.einsum('ij,jk,ik->',A,B,A,optimize=True)