NumPy 中实现二维数组与二维权重矩阵的向量化逐行加权求和

本文介绍如何使用 `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 表达数学意图，再根据性能分析决定是否优化。