
本文介绍如何使用 `np.einsum` 或广播机制,高效实现矩阵 `a` 每行分别与权重矩阵 `c` 各列做对应元素相乘后沿行方向(axis=0)求和,避免显式 python 循环,兼顾可读性与性能。
在科学计算中,常需对一个二维数组 a(形状为 (m, n))的每一行,分别应用一组权重(如来自多项式基或时间衰减因子),再按列求和得到结果向量。例如,给定:
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(0, 3)][::-1]) # shape: (3, 2)
# c = [[1. , 0.25],
# [1. , 0.5 ],
# [1. , 1. ]]目标是:对 c 的每一列(共 2 列),执行如下操作:
- 将 a 的第 i 行与 c[i, col] 相乘(标量乘整行);
- 对所有行的结果累加,得到长度为 n=3 的向量;
- 最终输出形状为 (n, c.shape[1]) = (3, 2) 的结果矩阵。
✅ 推荐解法:np.einsum(最清晰、高效)
利用爱因斯坦求和约定,明确指定索引关系:
result = np.einsum('ij,ik->jk', a, c)
print(result)
# [[32. 11. ]
# [50. 29. ]
# [54. 40.5]]其中 'ij,ik->jk' 含义为:
- i: 求和轴(a 的行索引,c 的行索引,二者对齐后求和);
- j: a 的列索引(保留为输出第 0 维);
- k: c 的列索引(保留为输出第 1 维)。
即:对每个 j,k,计算 ∑_i a[i,j] * c[i,k] —— 正是所需的逐列加权列和。
⚠️ 注意:np.einsum('ij,ik', a, c)(无箭头)默认对重复下标 i 求和,等价于 'ij,ik->jk',但显式写出更利于理解和调试。
✅ 替代解法:广播 + sum()(更直观,稍低效)
通过维度扩展实现广播:
result = (a[:, None, :] * c[:, :, None]).sum(axis=0) # a[:, None, :] → (3, 1, 3) # c[:, :, None] → (3, 2, 1) # 广播后形状:(3, 2, 3),再沿 axis=0 求和 → (2, 3) result = result.T # 转置为 (3, 2),与 einsum 一致
但注意:此写法输出为 (2, 3),需 .T 才匹配 einsum 的 (3, 2) 结构。若坚持 (3, 2) 输出,可改用:
result = (a.T[:, None] * c).sum(axis=2).T # 更紧凑的广播形式
? 关键对比总结: | 方法 | 可读性 | 性能 | 内存开销 | 推荐场景 | |--------------|--------|------|----------|------------------------| | np.einsum | ★★★★★ | ★★★★☆ | 低 | 首选,逻辑清晰、通用性强 | | 广播 + sum | ★★★☆☆ | ★★★☆☆ | 中高 | 理解广播机制时辅助验证 |
? 扩展提示:若后续需对 c 做归一化(如每列和为 1),可在 einsum 后直接处理;若 c 极大,可考虑分块计算或 numba 加速。始终优先用 einsum 表达数学意图,再根据性能分析决定是否优化。









