python - 块状整体矩阵与单行相乘

我有两个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的方法来实现这一目标？
谢谢

最佳答案

我们可以使用np.einsum-

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

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

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

关于python - 块状整体矩阵与单行相乘，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/53113941/