如何向量化实现两个numpy数组的成对求和？

我有两个numpy数组，结构如下：

x = [v1, v2, v3, ..., vm]
y = [w1, w2, w3, ..., wn]

其中vi, wj是长度为3的numpy数组。我想要对v和w进行成对求和，得到最终数组：

z = [v1+w1, v1+w2,...,v1+wn,v2+w1, ..., vi+wj, ..., vm+wn]

目前我用的简单方法是嵌套循环：

z = np.zeros ((m*n, 3))
for i in range(m):
    for j in range(n):
        z[n*i+j] = x[i] + y[j] 

但当m, n非常大时，这个方案效率太低。我知道scipy.spatial里有distance_matrix可以向量化计算成对距离，想问问有没有针对numpy数组的此类成对求和的向量化实现？

可以利用numpy的广播机制实现完全向量化的成对求和，不需要嵌套循环，效率会高很多。具体步骤如下：

1. 扩展x和y的维度，让它们的形状分别变成(m, 1, 3)和(1, n, 3)，这样numpy会自动广播维度，实现所有元素的成对相加；
2. 相加后得到形状为(m, n, 3)的数组；
3. 最后将结果reshape为(m*n, 3)，就是你需要的z数组。

对应的代码实现：

import numpy as np

# 示例数据
m, n = 2, 3
x = np.random.rand(m, 3)
y = np.random.rand(n, 3)

# 向量化成对求和
z = (x[:, np.newaxis, :] + y[np.newaxis, :, :]).reshape(-1, 3)

你也可以用np.expand_dims替代np.newaxis，写法更直观：

z = (np.expand_dims(x, axis=1) + np.expand_dims(y, axis=0)).reshape(-1, 3)

验证结果和循环方法的一致性：

# 循环方法生成z_loop
z_loop = np.zeros((m*n, 3))
for i in range(m):
    for j in range(n):
        z_loop[n*i+j] = x[i] + y[j]

# 检查是否相等（考虑浮点精度误差）
print(np.allclose(z, z_loop))  # 输出 True

这种向量化方式完全利用numpy的底层优化，避免了Python循环的开销，在m和n很大时性能提升非常明显。

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