Numpy整个矩阵乘以单行

Question

我有两个 numpy 矩阵，A 和 B：A = [[0,-1],     [-2,5],     [2,6]]B = [[5,3],     [2,4]]Dimension of A: N * KDimension of B: K * K我想要做的是对A 的每一行和整个 B进行 marix 乘法，然后再将它与A 的每一行相乘。然后总结一下。这是：N = A.shape[0]sum = 0for i in range(N):    s = A[i] @ B @ A[i]    print(s)    sum += sreturn sum这印470224而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)