如何向量化Pandas DataFrame的score函数计算以提升运行效率

我有一个包含val和target列的Pandas DataFrame，其中target为布尔类型，生成代码如下：

import random
import numpy as np
import pandas as pd

df = pd.DataFrame({'val': np.random.uniform(-1., 1., 1000), 
                   'target': random.choices([True, False], k=1000)})

需要针对多组不同的lo_lim和up_lim参数对，应用以下score函数：

def score(df, lo_lim, up_lim, alpha):
    df_out = df['target'].values[np.where((df['val']>up_lim) | (df['val']<lo_lim))[0]]    
    return df_out.sum()-alpha*(len(df_out)-df_out.sum())

当前采用循环遍历参数对的方式计算：

lo_lims = np.random.uniform(-1., -0.5, 100)
up_lims = np.random.uniform(0.5, 1.0, 100)

res = []
for i in range(100):
    res.append((lo_lims[i], up_lims[i], score(df, lo_lims[i], up_lims[i], 0.5)))

现在需要通过向量化处理一次性计算所有参数对的结果，大幅提升计算效率。

向量化实现方案 #

首先拆解score函数的数学逻辑，将其转化为可批量计算的形式：
原函数的计算式可整理为：

score = (1 + alpha) * 正样本数 - alpha * 总选中样本数

基于Numpy的广播机制，我们可以实现全量参数的并行计算：

1. 准备数组

将DataFrame的列转换为Numpy数组，并定义固定参数：

val_arr = df['val'].values  # 形状：(1000,)
target_arr = df['target'].values.astype(int)  # 布尔转int方便求和，形状：(1000,)
alpha = 0.5

2. 调整参数数组形状以支持广播

将lo_lims和up_lims转换为二维数组，以便和val_arr进行逐元素比较：

lo_arr = lo_lims[:, np.newaxis]  # 形状：(100, 1)
up_arr = up_lims[:, np.newaxis]  # 形状：(100, 1)

3. 生成掩码矩阵

计算每个参数对下，val满足val > up_lim或val < lo_lim的掩码：

# 掩码形状：(100, 1000)，每一行对应一组参数对的筛选结果
mask = (val_arr > up_arr) | (val_arr < lo_arr)

4. 批量计算样本数

计算每组参数对对应的总选中样本数和正样本数：

total_samples = mask.sum(axis=1)  # 形状：(100,)，每行求和得到总选中数
positive_samples = (mask * target_arr).sum(axis=1)  # 形状：(100,)，掩码与target相乘后求和得到正样本数

5. 批量计算score

代入整理后的公式计算所有参数对的score：

scores = (1 + alpha) * positive_samples - alpha * total_samples

6. 组合结果

将参数对和对应的score组合成最终结果：

# 转为Numpy数组形式
res_array = np.column_stack((lo_lims, up_lims, scores))
# 若需要列表元组的形式
res_list = list(map(tuple, res_array))

结果验证 #

可以对比向量化结果与原循环结果的一致性：

# 原循环结果
res_loop = []
for i in range(100):
    res_loop.append((lo_lims[i], up_lims[i], score(df, lo_lims[i], up_lims[i], 0.5)))

# 验证数值一致性
print(np.allclose(np.array(res_loop)[:, 2], scores))  # 输出True则说明结果一致

效率说明 #

Numpy的广播机制基于底层C实现，避免了Python循环的解释器开销。当参数对数量较大（如万级以上）时，向量化方案的速度会比循环快数十倍甚至上百倍。

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