如何用NumPy高效计算特定高维矩阵乘法（无Python循环）

输入变量整理 #

先把分散的n个矩阵和向量堆叠为统一的张量/矩阵，方便后续向量化操作：

• 矩阵A：形状(m, 3)，保持原格式即可
• 将n个(3,3)矩阵S₁~Sₙ堆叠为张量S，形状(n, 3, 3)，执行S = np.array([S₁, S₂, ..., Sₙ])即可实现
• 将n个(3,)向量v₁~vₙ堆叠为矩阵V，形状(n, 3)，执行V = np.array([v₁, v₂, ..., vₙ])即可实现

核心计算逻辑 #

利用NumPy的广播机制和向量化操作，完全避免Python循环：

import numpy as np

# 1. 计算每个Aᵢ与vⱼ的差值向量，得到形状为(m, n, 3)的张量diff
# A[:, None, :] 将A从(m,3)扩展为(m,1,3)，与V(n,3)广播后匹配为(m,n,3)
diff = A[:, None, :] - V

# 2. 计算差值向量与Sⱼ的矩阵乘法，得到形状为(m, n, 3)的张量temp
# diff(m,n,3) 与 S(n,3,3) 广播后做矩阵乘法，结果为(m,n,3)
temp = diff @ S

# 3. 将temp与diff对应元素相乘后求和，得到最终的(m,n)矩阵M
# 沿最后一个维度（长度为3）求和，每个位置得到标量结果
M = np.sum(temp * diff, axis=-1)

用Einstein求和简化写法 #

如果你熟悉Einstein求和符号，也可以用np.einsum一步完成计算，可读性更强：

diff = A[:, None, :] - V
M = np.einsum('ijk,jkl,ijl->ij', diff, S, diff)

这个表达式直接对应数学公式：对每个i,j，计算$\sum_{k,l} d_{ijk} S_{jkl} d_{ijl}$，其中$d_{ijk} = A_{ik} - V_{jk}$，和原式完全等价。

正确性验证示例 #

取小规模数据手动验证：比如m=1, n=1，A=[[1,2,3]]，S₁=np.eye(3)，v₁=[0,0,0]，则$M_{00}=(1,2,3) \cdot I \cdot (1,2,3)^T=14$，用上述代码计算会得到相同结果。

内容的提问来源于stack exchange，提问作者Cider