pandas按两列分组基于日期列计算4天滚动平均值方法

需求说明 #

需要实现的计算逻辑：

• 先按A、B两列对数据集分组
• 分组内基于日期列D，计算4天时间跨度的C列滚动平均值
• 原有实现的问题：传入整数4作为滚动窗口参数时，是按固定4条记录做计算，不考虑日期间隔，不符合时间窗口要求。例如当A=278、B=2、日期为2013-01-09时，距离上一条记录2013-01-06间隔3天，仅这两条记录落在4天窗口内，正确均值应为(2+3)/2=2.5。

原有错误实现代码：

df_time = pd.DataFrame({'A': [123, 123, 278, 278, 278, 123, 345, 278, 123,278, 278],
                        'B': [1, 2, 2, 2, 2, 1, 3, 2, 1, 2, 2],
                        'C': [0, 4, 2, 0, 4, 5, 3, 2, 1, 3, 2],
                       'D' : [pd.Timestamp('20130101'),
                                pd.Timestamp('20130102'),
                                pd.Timestamp('20130101'),
                                pd.Timestamp('20130102'),
                                pd.Timestamp('20130103'),
                                pd.Timestamp('20130103'),
                                pd.Timestamp('20130104'),
                                pd.Timestamp('20130104'),
                                pd.Timestamp('20130105'),
                                pd.Timestamp('20130106'),
                                pd.Timestamp('20130109')],
                       'rol_avg': [0, 2, 2, 1, 2, 5, 3, 3, 3, 2, 2.5] } )

df_time.sort_values(by='D')

sum_df = df_time.groupby(by=['A', 'B', 'D'])['C'].sum()
# 错误点：传入整数4，按4条记录滚动，而非4天时间窗口
rolling_average_series = sum_df.groupby(level=[0, 1]).rolling(4, min_periods=1).mean()

实现方案 #

核心修改点：将滚动窗口参数从整数4改为时间偏移字符串'4D'，同时保证分组内传入rolling方法的索引为日期时间类型，pandas会自动按时间跨度筛选窗口内的记录，不再按固定行数计算。

完整可运行代码：

import pandas as pd

df_time = pd.DataFrame({'A': [123, 123, 278, 278, 278, 123, 345, 278, 123,278, 278],
                        'B': [1, 2, 2, 2, 2, 1, 3, 2, 1, 2, 2],
                        'C': [0, 4, 2, 0, 4, 5, 3, 2, 1, 3, 2],
                       'D' : [pd.Timestamp('20130101'),
                                pd.Timestamp('20130102'),
                                pd.Timestamp('20130101'),
                                pd.Timestamp('20130102'),
                                pd.Timestamp('20130103'),
                                pd.Timestamp('20130103'),
                                pd.Timestamp('20130104'),
                                pd.Timestamp('20130104'),
                                pd.Timestamp('20130105'),
                                pd.Timestamp('20130106'),
                                pd.Timestamp('20130109')],
                       'rol_avg': [0, 2, 2, 1, 2, 5, 3, 3, 3, 2, 2.5] } )

# 第一步：聚合得到每个(A,B)分组下每日C列的总和，按日期排序
sum_df = df_time.groupby(by=['A', 'B', 'D'], as_index=False)['C'].sum().sort_values(by=['A','B','D'])

# 第二步：分组后设置日期列为时间索引，使用4D时间窗口计算滚动均值
rolling_result = sum_df.groupby(['A', 'B'], group_keys=False).apply(
    lambda group: group.set_index('D')['C'].rolling('4D', min_periods=1).mean()
).reset_index(name='calc_rol_avg')

# 合并结果得到最终输出
final_df = sum_df.merge(rolling_result, on=['A', 'B', 'D'])

结果验证 #

针对示例中A=278、B=2、D=2013-01-09的记录：

• 4天窗口的时间范围为2013-01-06至2013-01-09
• 窗口内仅包含2013-01-06的C值3、当日C值2
• 计算得到均值为(3+2)/2=2.5，和预期结果完全匹配。

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