百万行Pandas数据：双列分组下前序日期DPD均值计算优化

问题描述 #

现有如下结构的Pandas DataFrame：

import pandas as pd

df = pd.DataFrame([['John', 'A', '1/1/2017', '10'],
                   ['John', 'A', '2/2/2017', '15'],
                   ['John', 'A', '2/2/2017', '20'],
                   ['John', 'A', '3/3/2017', '30'],
                   ['Sue', 'B', '1/1/2017', '10'],
                   ['Sue', 'B', '2/2/2017', '15'],
                   ['Sue', 'B', '3/2/2017', '20'],
                   ['Sue', 'B', '3/3/2017', '7'],
                   ['Sue', 'B', '4/4/2017', '20']],
                  columns=['Customer', 'Group', 'Deposit_Date', 'DPD'])

需要新增PreviousMean列，该列表示同一Customer+Group分组下，截至当前Deposit_Date之前（不含当前日期）的DPD均值，无前置记录时为NaN。

原实现使用apply逐行计算，在百万行数据下速度极慢：

df['PreviousMean'] = df.apply(
    lambda x: df[(df.Customer == x.Customer) & 
                 (df.Group == x.Group) & 
                 (df.Deposit_Date < x.Deposit_Date)].DPD.mean(), 
axis=1)

高效实现方案 #

利用Pandas的矢量化分组聚合操作替代逐行循环，大幅提升处理效率：

步骤1：转换数据类型 #

先将日期列转为datetime类型，DPD转为数值类型，确保后续计算准确：

df['Deposit_Date'] = pd.to_datetime(df['Deposit_Date'])
df['DPD'] = df['DPD'].astype(int)

步骤2：按日期聚合分组数据 #

对每个Customer+Group+Deposit_Date组合，计算当日的DPD总和与记录数：

date_agg = df.groupby(['Customer', 'Group', 'Deposit_Date']).agg(
    total_dpd=('DPD', 'sum'),
    count=('DPD', 'size')
).reset_index()

步骤3：计算前置累计值与均值 #

在Customer+Group分组内，计算累计的DPD总和与记录数，通过shift获取截至前一日的累计值，最终计算前置均值：

# 分组计算累计总和与记录数
date_agg['cum_total'] = date_agg.groupby(['Customer', 'Group'])['total_dpd'].cumsum()
date_agg['cum_count'] = date_agg.groupby(['Customer', 'Group'])['count'].cumsum()

# 移位得到截至前一日的累计值
date_agg['prev_total'] = date_agg.groupby(['Customer', 'Group'])['cum_total'].shift(1)
date_agg['prev_count'] = date_agg.groupby(['Customer', 'Group'])['cum_count'].shift(1)

# 计算前置均值
date_agg['PreviousMean'] = date_agg['prev_total'] / date_agg['prev_count']

步骤4：合并结果回原DataFrame #

将计算好的PreviousMean合并到原数据中：

df = df.merge(
    date_agg[['Customer', 'Group', 'Deposit_Date', 'PreviousMean']],
    on=['Customer', 'Group', 'Deposit_Date'],
    how='left'
)

效果验证 #

最终输出与预期完全一致：

Customer Group Deposit_Date  DPD  PreviousMean
0     John     A   2017-01-01   10           NaN
1     John     A   2017-02-02   15          10.0
2     John     A   2017-02-02   20          10.0
3     John     A   2017-03-03   30          15.0
4      Sue     B   2017-01-01   10           NaN
5      Sue     B   2017-02-02   15          10.0
6      Sue     B   2017-03-02   20          12.5
7      Sue     B   2017-03-03    7          15.0
8      Sue     B   2017-04-04   20          13.0

效率说明 #

该方案全程使用Pandas矢量化操作，避免了apply的逐行循环，时间复杂度从O(n²)降至O(n log n)，百万行数据的处理速度可提升数十倍甚至上百倍。

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