求Pandas大DataFrame分组日期区间内行数量的向量化统计方案

需求说明 #

给定如下Pandas DataFrame，需按Name分组，统计每条记录的Date往前推3个月内该Name的出现次数，同时生成对应的Date_3_Months_Ago字段。由于数据集规模极大，需避免低效的apply逐行处理或循环，采用向量化方案实现。

原始DataFrame #

Name    Date
Alice   2023-01-01
Alice   2023-02-01
Bob     2023-02-15
Alice   2023-03-01
Bob     2023-03-20
Alice   2023-04-01

预期输出 #

Name       Date Date_3_Months_Ago  N_Past_3_Months
0   Alice 2023-01-01        2022-10-01                1
1   Alice 2023-02-01        2022-11-01                2
3   Alice 2023-03-01        2022-12-01                3
5   Alice 2023-04-01        2023-01-01                4
2     Bob 2023-02-15        2022-11-15                1
4     Bob 2023-03-20        2022-12-20                2

向量化解决方案 #

核心思路 #

利用Pandas分组操作结合searchsorted向量化函数，基于二分查找快速定位符合时间范围的记录位置，避免逐行遍历，实现O(n log n)的时间复杂度。

代码实现 #

import pandas as pd

# 构造原始数据（实际场景可直接读取现有DataFrame）
df = pd.DataFrame({
    'Name': ['Alice', 'Alice', 'Bob', 'Alice', 'Bob', 'Alice'],
    'Date': ['2023-01-01', '2023-02-01', '2023-02-15', '2023-03-01', '2023-03-20', '2023-04-01']
})

# 1. 将Date列转换为datetime类型，确保时间计算正确性
df['Date'] = pd.to_datetime(df['Date'])

# 2. 按Name和Date排序，保证组内时间有序（searchsorted依赖有序序列）
df_sorted = df.sort_values(['Name', 'Date']).reset_index(drop=True)

# 3. 计算每条记录对应的3个月前日期
df_sorted['Date_3_Months_Ago'] = df_sorted['Date'] - pd.DateOffset(months=3)

# 4. 组内向量化统计符合条件的次数
def calc_past_count(group):
    dates = group['Date'].values
    past_dates = group['Date_3_Months_Ago'].values
    # 用二分查找找到每个past_date在组内日期序列中的起始位置
    start_indices = dates.searchsorted(past_dates, side='left')
    # 当前行位置减去起始位置再加1，即为近3个月内的出现次数（包含当前记录）
    group['N_Past_3_Months'] = group.index.get_level_values(1) - start_indices + 1
    return group

# 分组应用计算函数（组内为向量化操作，效率远高于逐行处理）
result = df_sorted.groupby('Name', group_keys=False).apply(calc_past_count)

# 可选：恢复原始DataFrame的行顺序（若需与预期输出的行索引一致）
result = result.loc[df.index].reset_index(drop=True)

print(result)

方案优势 #

• 效率极高：searchsorted基于二分查找，组内处理时间为O(m log m)（m为组内行数），整体时间复杂度O(n log n)，相比逐行apply的O(n²)，在大数据集下性能提升显著。
• 内存友好：避免了全量广播布尔索引带来的内存开销，适合处理千万级以上规模的数据集。

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