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进行 marix 乘法，然后再将它与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这正是我想要的。在没有 for 循环的情况下，是否有任何优雅或 Pythonic 的方式来实现这一点？谢谢

Answer 1

我们可以使用np.einsum——np.einsum('ij,jk,ik->',A,B,A)并使用optimize标志进一步优化它-np.einsum('ij,jk,ik->',A,B,A,optimize=True)