基于类别与滑动日期范围高效计算条件求和的最优方法

问题说明 #

给定结构如下的Pandas DataFrame：

需要给每行算出同类别、且日期在该行日期前后14天范围内的所有行X-Value的总和。数据集有50万行，当前用的逐行遍历全表的方法速度慢到离谱，按类别+日历月份分组的方式又满足不了每行动态滑动窗口的需求。当前的低效代码如下：

def cohort_value(row):
    row_date = row['date']
    window_start_date = row_date - timedelta(days=14)
    window_end_date = row_date + timedelta(days=14)
    
    cohort_value = df[(df['date'] > start_date) & (df['date'] < end_date)]['value'].sum()
    
    return  cohort_value

df.apply(cohort_value, axis=1)

高效解决办法 #

针对50万行的规模，别再用逐行遍历了，直接用分组+滚动窗口或者二分查找定位窗口范围的方法，速度能提几十上百倍，下面是两种可行方案：

方案一：分组排序后用时间滚动窗口（适合日期比较连续的场景） #

1. 先按Category分组，每组内按Date排序
2. 用Pandas的rolling窗口，基于时间偏移设置28天的窗口（前后14天），直接计算总和
3. Pandas 1.3及以上版本支持closed='both'，能同时包含窗口两端的日期

import pandas as pd
from datetime import timedelta

# 先把Date列转成datetime类型，避免后续出错
df['Date'] = pd.to_datetime(df['Date'])

# 按类别分组，每组内按日期排序
df_sorted = df.sort_values(['Category', 'Date']).reset_index(drop=True)

# 对每组应用时间滚动窗口，计算前后14天的X-Value总和
df_sorted['Rolling_Sum'] = df_sorted.groupby('Category').apply(
    lambda g: g.set_index('Date')['X-Value'].rolling(
        window='28D',  # 前后14天加起来28天的时间窗口
        closed='both'
    ).sum()
).reset_index(level=0, drop=True)

# 如果需要恢复原数据的顺序，就用merge合并回去
df = df.merge(df_sorted[['Date', 'Category', 'Rolling_Sum']], on=['Date', 'Category'], how='left')

方案二：用二分查找定位窗口范围（适合日期稀疏的场景） #

如果你的日期分布很零散，滚动窗口可能不太适配，那就用searchsorted做二分查找，快速定位每个日期对应的窗口索引范围，再求和：

import pandas as pd
from datetime import timedelta

# 同样先确保Date是datetime类型
df['Date'] = pd.to_datetime(df['Date'])

# 按类别分组并排序
df_sorted = df.sort_values(['Category', 'Date']).reset_index(drop=True)
result_groups = []

for cat, group in df_sorted.groupby('Category'):
    dates = group['Date'].values
    x_vals = group['X-Value'].values
    # 对组内每一行，用二分查找找到窗口的左右边界索引
    for idx, current_date in enumerate(dates):
        # 找到大于等于current_date-14天的第一个索引
        left_idx = dates.searchsorted(current_date - timedelta(days=14))
        # 找到大于current_date+14天的第一个索引，取前一位就是窗口右边界
        right_idx = dates.searchsorted(current_date + timedelta(days=14), side='right')
        # 计算该范围内的X-Value总和
        group.loc[group.index[idx], 'Rolling_Sum'] = x_vals[left_idx:right_idx].sum()
    result_groups.append(group)

# 合并所有组，再和原表匹配恢复顺序
df = pd.concat(result_groups).merge(df, on=['Date', 'Category', 'X-Value'], how='right')

为什么这两种方法快？ #

• 方案一用了Pandas底层优化的滚动窗口，是C级别的计算，比Python逐行循环快N倍
• 方案二用二分查找（O(log n)复杂度）代替全表遍历（O(n)），整个流程的时间复杂度是O(n log n)，50万行完全能轻松处理
• 两种方法都把计算限制在同类别组内，不用每次都扫全表，计算量直接砍了大半

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