NumPy中如何无for循环向量化实现两数组相加得到[n,m,3]数组

实现思路 #

直接利用NumPy的广播机制完成向量化计算，不需要编写任何Python层面的for循环：
广播规则要求参与运算的数组维度对齐，我们只需要给形状为(n,3)的数组a插入一个长度为1的新轴，把形状调整为(n,1,3)，和形状为(m,3)的数组b做加法时，NumPy会自动完成维度适配：

1. 自动给b在最前端补长度为1的轴，形状变为(1,m,3)
2. 两个数组中长度为1的轴会自动广播复制，匹配到n和m的长度
   最终直接得到形状为(n,m,3)的结果数组，天然满足c[i,j,:] = a[i,:] + b[j,:]的计算规则。

代码示例 #

import numpy as np

# 测试数据，可替换为实际使用的a、b数组
n = 3
m = 4
a = np.array([[1,2,3], [4,5,6], [7,8,9]])  # 形状(3,3)
b = np.array([[10,20,30], [40,50,60], [70,80,90], [100,110,120]])  # 形状(4,3)

# 核心计算逻辑：np.newaxis等价于None，也可简写为a[:, None] + b
c = a[:, np.newaxis] + b

# 输出c的形状为(3,4,3)，和预期形状一致
print(c.shape)

可以随机抽取索引验证结果正确性，比如c[0,1]的计算结果为[41,52,63]，和手动计算a[0] + b[1] = [1,2,3] + [40,50,60]的结果完全一致。

性能优势 #

这种实现的所有循环逻辑都在NumPy底层C层执行，相比Python层手写两层for循环遍历i、j索引的实现，当n、m规模较大时（比如达到万级），计算速度可以提升两个数量级以上。

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