
本文详解如何用 `np.einsum` 或广播机制,对矩阵 a 的每行分别按系数矩阵 c 的各列进行加权并沿行方向求和,避免显式 python 循环,实现高效向量化计算。
在科学计算中,常需对一个二维数组 a(形状为 (m, n))的每一行,分别乘以一个系数向量(如 c 的某一列),再沿行维度(axis=0)求和,得到长度为 n 的结果向量。当系数本身构成二维数组 c(形状为 (m, k))时,我们希望同时对 c 的每一列(共 k 列)执行该操作,最终输出形状为 (n, k) 的结果矩阵——即:第 j 列结果 = sum(a[i, :] * c[i, j] for i in range(m))。
直接使用循环虽直观,但性能差;而 a * c[:, None] 仅适用于 c 为一维的情形。针对二维 c,推荐两种向量化方案:
✅ 方案一:np.einsum(推荐,语义清晰、性能优)
利用爱因斯坦求和约定,明确指定索引关系:
- a 索引为 ij(i: 行,j: 列)
- c 索引为 ik(i: 行,k: 系数列)
- 输出省略求和下标 i,保留 j, k → 结果形状为 (n, k)
import numpy as np
a = np.array([[20, 12, 6],
[12, 24, 18],
[ 0, 14, 30]])
b = np.array([1, 0.5])
c = np.array([b ** i for i in range(3)][::-1]) # shape: (3, 2)
result = np.einsum('ij,ik->jk', a, c) # 注意显式写出输出下标 'jk'
print(result)
# [[32. 11. ]
# [50. 29. ]
# [54. 40.5]]✅ 输出为 (n, k) = (3, 2),符合预期:每列对应 c 的一列权重下的加权和。
✅ 方案二:广播 + 轴向求和(兼容性强,稍低效)
通过维度扩展实现广播:
- a[:, None] → (3, 1, 3)
- c[:, :, None] → (3, 2, 1)
- 广播后形状为 (3, 2, 3),再对 axis=0 求和 → (2, 3)
- 最后转置得 (3, 2)(或直接 sum(0).T)
result_broadcast = (a[:, None] * c[:, :, None]).sum(axis=0).T # 等价于:(a[:, None] * c.T[None, :, :]).sum(0)
⚠️ 注意:此方式内存占用更高(生成中间三维数组),在大数据量时不如 einsum 节省内存。
? 验证逻辑(以 c 第二列 [0.25, 0.5, 1.0] 为例):
# 手动计算第二列(索引 1): col1 = c[:, 1] # [0.25, 0.5, 1.0] manual = (a.T * col1).sum(axis=1) # [32. , 50. , 54. ] # 对应 result[:, 1] → 完全一致
✅ 总结
- 首选 np.einsum('ij,ik->jk', a, c):语义明确、性能好、不易出错;
- 广播方案可作为理解维度变换的辅助手段,但生产环境建议优先使用 einsum;
- 所有操作均完全向量化,无需 Python 循环,充分利用 NumPy 底层优化。










