如何加速DataFrame循环？90万条数据的向量化优化方法

问题描述 #

需要处理90万条数据，希望加速DataFrame中的循环操作。需求是为DataFrame新增名为Count的列，统计同一player中score小于当前行target score的次数（每行的target score不同）。

输入示例 #

index Nom player  score  target score
0      0      felix      3            10
1      1      felix      8             7
2      2       theo      4             5
3      3    patrick     12             6
4      4     sophie      7             6
5      5     sophie      3             6
6      6      felix      2             4
7      7      felix      2             2
8      8      felix      2             3

预期输出 #

index Nom player  score  target score  Count
0      0      felix      3            10      5
1      1      felix      8             7      4
2      2       theo      4             5      1
3      3    patrick     12             6      0
4      4     sophie      7             6      1
5      5     sophie      3             6      1
6      6      felix      2             4      4
7      7      felix      2             2      0
8      8      felix      2             3      3

当前低效代码 #

df2 = df.copy()
df2['Count']= [np.count_nonzero((df.values[:,1] == row[2] )& ( df.values[:,2] < row[4]) )    for row in df.itertuples()]
print(df2)

询问是否可以通过向量化优化该计算，以及具体实现方法。

向量化优化方案 #

当然可以通过向量化操作大幅提升效率，核心思路是分组排序+二分查找，将原O(n²)的逐行循环优化为O(n log k)的高效计算（k为每个player组的大小），完全适配90万条数据的场景。

实现方式1：分组预排序 + 逐行查找 #

import pandas as pd
import numpy as np

# 预先生成每个player对应的排序后score数组
sorted_score_map = df.groupby('player')['score'].apply(np.sort)

# 定义查找函数：在对应player的排序数组中，用二分查找统计小于target的元素数
def calc_count(row):
    return np.searchsorted(sorted_score_map[row['player']], row['target score'], side='left')

# 批量计算Count列
df['Count'] = df.apply(calc_count, axis=1)

实现方式2：分组批量处理（更高效） #

这种方式直接在分组内批量处理所有行，避免逐行apply的开销：

import pandas as pd
import numpy as np

# 分组后对每个组内的score排序，再批量用二分查找匹配target score
counts = df.groupby('player').apply(
    lambda group: np.searchsorted(np.sort(group['score']), group['target score'], side='left')
)

# 将分组结果展开并赋值回原DataFrame
df['Count'] = counts.explode().values

效率说明 #

• 原代码的逐行循环会对每行遍历整个DataFrame，数据量越大效率越低；优化后的方案利用numpy的二分查找（底层C实现）和pandas的分组操作，将时间复杂度从O(n²)降到O(n log k)，90万条数据的处理速度会有数量级的提升。
• np.searchsorted的side='left'参数确保返回第一个大于等于目标值的位置，这个位置的索引恰好是数组中小于目标值的元素总数，完全符合需求。

内容的提问来源于stack exchange，提问作者Félix Le Besnerais