Pandas按number分组统计小于当前行日期的记录数及性能优化

问题分析 #

原代码使用apply逐行过滤整个DataFrame，时间复杂度为O(n²)——每一行都要对全表进行匹配过滤，当数据量达到30万行时，重复计算会导致计算量急剧膨胀，最终耗时长达2小时。

优化方案 #

以下两种方法均利用Pandas的向量化分组操作，将时间复杂度降至O(n log n)，能大幅提升处理速度：

方法1：使用groupby + transform + rank #

通过分组排序后计算排名，直接得到每个日期对应的“小于当前日期”的记录数：

import pandas as pd

# 初始化数据
df = pd.DataFrame({
    'id': [1,2,3,4,5,6],
    'number':['a', 'a', 'a', 'a', 'b', 'c'], 
    'date' : ['2023-06-20', '2022-06-20', '2021-06-20', '2021-07-20', '2022-06-20', '2023-06-20']
})

# 将日期列转换为datetime类型（避免字符串比较的潜在问题，更规范）
df['date'] = pd.to_datetime(df['date'])

# 分组后计算日期排名，rank(method='min')确保相同日期获得相同排名，减1得到小于当前日期的记录数
df['result'] = df.groupby('number')['date'].transform(
    lambda x: x.rank(method='min', ascending=True) - 1
)

print(df)

方法2：使用groupby + transform + searchsorted #

利用二分查找快速定位小于当前日期的记录数量，性能更优：

import pandas as pd

# 初始化数据
df = pd.DataFrame({
    'id': [1,2,3,4,5,6],
    'number':['a', 'a', 'a', 'a', 'b', 'c'], 
    'date' : ['2023-06-20', '2022-06-20', '2021-06-20', '2021-07-20', '2022-06-20', '2023-06-20']
})

df['date'] = pd.to_datetime(df['date'])

def count_prev_records(group):
    sorted_dates = group.sort_values().reset_index(drop=True)
    return group.apply(lambda x: sorted_dates.searchsorted(x))

df['result'] = df.groupby('number')['date'].transform(count_prev_records)

print(df)

结果验证 #

两种方法输出的result列均与原apply方法一致：

id number       date  result
0   1      a 2023-06-20     3.0
1   2      a 2022-06-20     2.0
2   3      a 2021-06-20     0.0
3   4      a 2021-07-20     1.0
4   5      b 2022-06-20     0.0
5   6      c 2023-06-20     0.0

性能对比 #

• 原apply方法：30万行数据耗时约2小时
• 优化方法：30万行数据仅需几秒到几十秒（具体取决于硬件配置）

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