Pandas按类别分组计算当前月与上月spendings列滚动均值的实现问题

需求说明 #

现有分属不同类别的跨月度时间序列数据，需要按类别分组，对spendings列计算当前月及存在的上月所有数据的滚动均值。

示例数据 #

import pandas as pd
from io import StringIO

data = StringIO(
"""
date          spendings  category
2020-01-01    10         A
2020-01-01    20         A
2020-01-01    15         A
2020-02-01    10         B
2020-02-01    10         B
2020-02-01    14         A
2020-02-01    19         A
2020-03-01    50         A
2020-04-01    40         A
"""
)

df = pd.read_csv(data, sep="\s+", parse_dates=True, index_col="date")

预期计算规则 #

• 类别A在2020-02月对应的均值为(10+20+15+14+19)/5 = 15.6
• 类别A在2020-03月对应的均值为(14+19+50)/3 ≈27.666667

原有实现的问题 #

原尝试使用df.groupby('category').rolling('30D').spendings.mean()存在两个缺陷：

• 仅计算30天内当前行及之前行的均值，未包含当月所有其他行的数值
• 30D参数不等同于1个自然月，会出现跨自然月统计错误的情况

正确实现方案 #

核心思路是先按类别+月份聚合得到每个月的总支出和记录数，再按类别做2个周期的滚动统计，用总支出除以总记录数得到均值：

# 1. 按类别分组后做月度重采样，统计每个月的总支出和记录条数
monthly_agg = df.groupby('category').resample('M').agg(
    total_spend=('spendings', 'sum'),
    count=('spendings', 'count')
).reset_index(level=0)

# 2. 按类别分组，滚动2个周期（覆盖当月+上月）计算累计总支出和累计记录数
rolling_agg = monthly_agg.groupby('category')[['total_spend', 'count']].rolling(2, min_periods=1).sum()

# 3. 计算滚动均值，调整索引顺序匹配预期格式
result = (rolling_agg['total_spend'] / rolling_agg['count']).swaplevel().sort_index()

输出结果 #

category  date      
A         2020-01-31    15.000000
          2020-02-29    15.600000
          2020-03-31    27.666667
          2020-04-30    45.000000
B         2020-02-29    10.000000
dtype: float64

如果需要将日期统一为当月1号，可额外添加一步日期调整：

result = result.reset_index()
result['date'] = result['date'] - pd.offsets.MonthBegin(1)
result = result.set_index(['category', 'date'])

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