Python/Numba中不同形状矩阵快速相加的最优方法

直接利用NumPy的广播机制就能实现无循环的高效计算，底层由优化过的C代码执行，大尺寸场景下性能远优于Python循环甚至Numba加速的循环。

核心思路 #

通过给两个矩阵扩展维度，让它们的形状匹配广播规则：

• 将形状为(K,T)的矩阵a扩展为(K,T,1)（在最后一维加一个长度为1的维度）
• 将形状为(K,N)的矩阵b扩展为(K,1,N)（在中间一维加一个长度为1的维度）

此时两个矩阵会自动广播到(K,T,N)的形状，对应位置相加即可得到目标结果。

代码实现 #

import numpy as np

# 生成示例数据
K, T, N = 100, 200, 300
a = np.random.rand(K, T)
b = np.random.rand(K, N)

# 方法1：用索引扩展维度
result = a[..., np.newaxis] + b[:, np.newaxis, :]

# 方法2：用np.expand_dims更直观
result = np.expand_dims(a, axis=2) + np.expand_dims(b, axis=1)

性能优势 #

NumPy的广播操作会直接调用底层优化的BLAS/LAPACK库，能充分利用CPU的向量指令集（如AVX、SSE），内存访问模式更连续，避免了Python循环的解释器开销。在K、T、N较大（比如均超过1000）的场景下，这种方案的运行速度通常是Numba三重循环的数倍。

结果验证 #

最终结果的每个元素result[k, t, n]严格等于a[k, t] + b[k, n]，和你之前的循环实现逻辑完全一致。

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