如何在Python/Numpy中高效地为笛卡尔积元素应用函数？

我懂你现在的困扰：单变量函数用map配合lambda处理起来很顺手，但遇到像def fun(x,y): return 2*x+y这类多变量函数，要把它应用到x、y的笛卡尔积每个元素时，map好像没法直接用，而且当前的替代方法还慢得让人着急。其实有几个高效的方案能解决这个问题，甚至比你之前的方法快很多，下面给你详细说说：

方法一：结合itertools.product和map实现 #

其实map是可以处理多变量场景的，关键是要把笛卡尔积的每个元素（元组形式）拆成函数的参数。你可以用itertools.product生成x和y的笛卡尔积，再通过lambda把元组拆包传给多变量函数：

import numpy as np
from itertools import product

def fun(x, y):
    return 2*x + y

xpts = np.linspace(0, 1, 5)
ypts = np.linspace(0, 2, 5)  # 这里y的范围可以根据你的需求修改

# 生成x和y的笛卡尔积，每个元素是(x,y)元组
cartesian_pairs = product(xpts, ypts)
# 用lambda拆包元组，传给fun函数
result = np.array(list(map(lambda pair: fun(*pair), cartesian_pairs)))

不过这种方法在数据量很大时，效率不如numpy原生的向量化操作，下面这个方法才是更推荐的。

方法二：numpy广播机制（速度最快，最推荐） #

numpy的核心优势就是向量化计算，完全不用循环或map，靠广播机制就能轻松处理笛卡尔积的计算场景，而且速度拉满：

import numpy as np

def fun(x, y):
    return 2*x + y

xpts = np.linspace(0, 1, 5)
ypts = np.linspace(0, 2, 5)

# 把x转成列向量，y转成行向量，利用广播自动生成笛卡尔积的计算结果
x_col = xpts[:, np.newaxis]  # 形状变为(5,1)
y_row = ypts[np.newaxis, :]  # 形状变为(1,5)
# 直接传入函数，numpy会自动广播计算所有(x,y)组合
result_matrix = fun(x_col, y_row)
# 如果需要一维数组形式的结果，直接扁平化即可
result_flat = result_matrix.flatten()

这种方法是numpy底层优化过的，数据量越大，和循环/map的速度差距越明显，绝对是处理这类问题的首选。

方法三：用numpy.vectorize快速适配现有函数 #

如果你不想改动原函数的结构，也可以用numpy.vectorize把多变量函数包装成能处理数组的向量化函数，虽然是“伪向量化”（底层还是循环），但胜在方便：

import numpy as np

def fun(x, y):
    return 2*x + y

xpts = np.linspace(0, 1, 5)
ypts = np.linspace(0, 2, 5)

# 包装函数，让它能处理数组输入
vec_fun = np.vectorize(fun)
# 生成x和y的网格矩阵
X, Y = np.meshgrid(xpts, ypts)
# 直接传入网格矩阵计算
result_matrix = vec_fun(X, Y)
result_flat = result_matrix.flatten()

这个方法适合快速改造现有代码，不用大动干戈，但速度不如方法二的纯广播操作。

总结一下：追求极致速度选方法二的numpy广播；一定要用map的话选方法一；想快速适配现有函数就用方法三。

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