如何加速二维数组间的逐元素函数运算（替代双重循环）

当你手里有一个只能处理单个浮点值的复杂函数（无法修改），现在要对两个同维度的二维NumPy数组$X_{n \times m}$、$Y_{n \times m}$的每对对应元素$x_{ij}$、$y_{ij}$执行这个函数，不想用低效的双重循环的话，直接上几种可行的提速方案：

方案1：用NumPy的vectorize快速封装 #

np.vectorize可以把处理单个元素的函数转换成能直接接收数组输入的函数，写法简洁，虽然底层还是循环，但代码可读性大幅提升，适合快速改造现有逻辑：

import numpy as np

def func(x, y):
    return x + y  # 替换为你的复杂函数

X = np.random.rand(100, 100)
Y = np.random.rand(100, 100)

# 封装函数
vectorized_func = np.vectorize(func)
# 直接对数组操作，无需循环
Z = vectorized_func(X, Y)

方案2：扁平化数组+列表推导 #

把二维数组拉成一维，用列表推导逐元素处理后再重塑回原形状，这种方式比嵌套循环更高效，代码也更简洁：

import numpy as np

def func(x, y):
    return x + y

X = np.random.rand(100, 100)
Y = np.random.rand(100, 100)

# 扁平化数组，逐对处理后重塑形状
Z_flat = [func(x, y) for x, y in zip(X.flat, Y.flat)]
Z = np.array(Z_flat).reshape(X.shape)

方案3：用Numba编译加速（极致性能） #

如果追求极致速度，用Numba的JIT装饰器编译你的函数，它会把Python代码转换成机器码，执行效率接近C语言，而且不需要修改原函数的逻辑：

import numpy as np
from numba import jit

# 用jit装饰器编译函数，nopython=True强制生成纯机器码，性能最优
@jit(nopython=True)
def func(x, y):
    return x + y

X = np.random.rand(100, 100)
Y = np.random.rand(100, 100)

# 直接对数组调用编译后的函数，Numba自动处理逐元素运算
Z = func(X, Y)

注意事项 #

• np.vectorize适合快速实现，性能提升有限，本质是Python循环的封装；
• 列表推导比嵌套循环快，但仍属于Python层面的循环，适合小体量数组；
• Numba是性能最优的选择，但需要额外安装库：pip install numba。

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