矩阵向量运算的Python向量化

Question

我有一个形状为（2,2，N）的矩阵A和一个形状为（2，N）的矩阵V

我要向量化以下内容：

F = np.zeros(N)
for k in xrange(N):
    F[k] = np.dot( A[:,:,k], V[:,k] ).sum()

可以通过tensordot或任何其他numpy函数完成此操作而无需显式循环吗？

Answer 1

使用np.einsum-

F = np.einsum('ijk,jk->k',A,V)

我们可以通过将optimize标志（检查文档）设置为True来对其进行进一步优化。