Pandas按A、C列分组的条件滚动平均值计算实现问题

更新说明 #

此前的同类解决方案未将A列与C列纳入分组条件考量。 #

需求说明 #

现有一个包含3列的DataFrame，需要新增2个计算列，具体规则如下：

• 按A、C列分组，对B列计算滚动2期平均值，取值为同组内当前行与前一行的B值均值
• 按A、C列分组，对B列计算滚动2期平均值，取值为同组内当前行之前的2行B值均值，无对应前序行时取值为0
  第二个计算项可使用日期或序列字段作为滚动计算的排序依据。

示例数据 #

import pandas as pd
df = pd.DataFrame({'A': ['t1', 't1', 't1', 't1', 't2', 't2', 't2', 't2','t1'],
                   'B': [100, 104, 108, 110, 102, 110, 98, 100, 200],
                   'C': ['h', 'a', 'a', 'a', 'a', 'h', 'h', 'h','h'],
                  'expected1': [100, 104, 106, 109, 102, 110, 104, 99, 150],
                  'expected2': [0, 0, 104, 106, 0, 0, 110, 104, 100]}, columns=['A', 'B', 'C','expected1','expected2'])

实现代码 #

# 计算第一个指标：包含当前行的2期滚动均值，最小统计周期为1
df['calc1'] = df.groupby(['A', 'C'], group_keys=False)['B'].rolling(window=2, min_periods=1).mean()

# 计算第二个指标：仅包含当前行之前的2期滚动均值，最小统计周期为2，空值填0
df['calc2'] = df.groupby(['A', 'C'], group_keys=False)['B'].apply(
    lambda x: x.shift(1).rolling(window=2, min_periods=2).mean()
).fillna(0)

结果验证 #

执行后calc1与expected1完全一致，calc2与expected2完全一致。如果原始数据需要按指定字段排序后再计算，只需要在分组前用sort_values()对排序字段做排序即可。

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