Numpy：对两个不同长度数组逐配对应用函数 替代循环实现

实现逻辑 #

利用NumPy广播机制自动生成两个数组所有元素的配对差值，完全替换双层循环：

• 假设X1长度为m，X2长度为n，将X1扩展为(m, 1)维度的二维数组
• 扩展后的X1与一维X2直接做差，会自动广播生成(m, n)的差值矩阵，矩阵中[i,j]位置的元素恰好对应你原来的X1[i] - X2[j]
• 差值矩阵直接传入计算方法即可得到最终结果

代码示例 #

情况1：my_method支持NumPy数组输入（性能最优） #

如果你的自定义方法本身可以直接对整个数组做运算（比如是NumPy内置的数学运算，或者自定义时用了NumPy的向量化运算），直接用以下代码，性能比双层循环提升数十到数百倍：

import numpy as np

# X1[:, np.newaxis] 实现维度扩展，触发广播机制
result = my_method(X1[:, np.newaxis] - X2)

小提示：X1[:, np.newaxis]也可以简写为X1[:, None]，二者效果完全相同

情况2：my_method仅支持单元素输入 #

如果你的自定义方法暂时只能接收单个数值输入，可以用np.vectorize做语法层面的简化，注意该方法本质还是循环，仅简化代码写法，性能提升有限，最优方案是改造my_method适配数组运算：

import numpy as np

# 包装自定义方法为支持数组输入的形式
vec_method = np.vectorize(my_method)
result = vec_method(X1[:, np.newaxis] - X2)

效果验证 #

你可以用以下测试用例验证输出和你原循环代码完全一致：

# 示例自定义方法：求平方
def my_method(x):
    return x ** 2

X1 = np.array([1, 2, 3])
X2 = np.array([4, 5])

# 向量化输出结果：
# [[ 9 16]
#  [ 4  9]
#  [ 1  4]]

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