如何用Numpy无循环实现(a,d)与(b,d)数组相加得到(a*b,d)形状的结果

最优方案：利用广播机制直接实现 #

不需要调用np.einsum或np.meshgrid，用numpy自带的广播规则就能一行实现，效率和可读性都最高，输出结果和你写的for循环完全一致：

import numpy as np

a = 2
b = 2
d = 3

matrix1 = np.random.rand(a,d)
matrix2 = np.random.rand(b,d)

# 无循环实现，输出形状(a*b, d)
result = (matrix1[:, np.newaxis, :] + matrix2).reshape(-1, d)

原理说明 #

给matrix1插入一个长度为1的中间维度后，形状从(a,d)变为(a,1,d)，和形状为(b,d)的matrix2执行加法时，numpy会自动触发广播机制，生成形状为(a,b,d)的中间结果，最后将前两个维度展平就得到你需要的输出格式，元素顺序和原for循环的遍历顺序完全匹配。

可选方案：使用你提到的np.einsum或np.meshgrid实现 #

这两个方法也能实现需求，但简洁度和性能都不如广播方案，仅做参考：

• 使用np.einsum实现：

# 本质还是借助广播逻辑，einsum在这里仅用于维度扩展，没有额外优势
result_einsum = (np.einsum('ik,j->ijk', matrix1, np.ones(b)) + matrix2).reshape(-1, d)

• 使用np.meshgrid实现：

# 先生成所有组合的索引，再按索引取值相加
idx1, idx2 = np.meshgrid(np.arange(a), np.arange(b), indexing='ij')
result_mesh = matrix1[idx1.flatten()] + matrix2[idx2.flatten()]

你可以通过np.allclose(result, 原for循环输出)验证所有方案的结果完全一致。

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