按组计算扩展均值时忽略前n行（含当前行）的实现方法

现有如下DataFrame：

a  b
week       
8     10  9
9      3  8
9      5  5
9      7  2
10     1  3
9      4  4
9      2  6

需求 #

按索引（week）分组计算各列的扩展均值，但忽略当前行及之前的n行（此处n=2），同时保留原DataFrame的行顺序。

预期输出（n=2时） #

a    b
week          
8     NaN  NaN
9     NaN  NaN
9     NaN  NaN
9     3.0  8.0
10    NaN  NaN
9     5.0  5.0
9     5.0  5.0

说明：

• 第四行（a列值3）：属于week=9分组，排除当前行及前一行（第三、四行）后，仅保留第二行的a值3，均值为3.0
• 最后一行（b列值5）：属于week=9分组，排除当前行及前一行（第六、七行）后，保留第二、三、四行的b值8、5、2，均值为(8+5+2)/3=5.0

尝试过的方法及问题 #

原按组计算扩展均值的代码：

(df.reset_index()
   .groupby('week')
   .expanding()
   .mean()
   .sort_index(level=1)
   .reset_index(level=1, drop=True))

尝试用shift(2)实现忽略前2行，但输出不符合预期：

(df.shift(2)
   .reset_index()
   .groupby('week')
   .expanding()
   .mean()
   .sort_index(level=1)
   .reset_index(level=1, drop=True))

输出结果：

a         b
week                     
8           NaN       NaN
9           NaN       NaN
9     10.000000  9.000000
9      6.500000  8.500000
10     5.000000  5.000000
9      6.666667  6.333333
9      5.250000  5.500000

解决方案 #

要实现需求，需要先为每个分组内的行标记序号，然后计算每个行对应的分组内序号小于当前序号-n的行的均值，最后按原顺序拼接结果。

代码实现：

import pandas as pd

df = pd.DataFrame(
    {'a': [10, 3, 5, 7, 1, 4, 2],
     'b': [9, 8, 5, 2, 3, 4, 6]},
    index=pd.Index([8,9,9,9,10,9,9], name='week')
)

n = 2

# 重置索引并为每个分组添加行号
df_with_idx = df.reset_index().assign(group_row_num=lambda x: x.groupby('week').cumcount())

# 定义函数：对每个分组，计算每行对应的符合条件的均值
def calc_lagged_expanding_mean(group):
    means = []
    for idx in group['group_row_num']:
        mask = group['group_row_num'] < (idx - n)
        if mask.any():
            means.append(group.loc[mask, ['a', 'b']].mean())
        else:
            means.append(pd.Series([pd.NA, pd.NA], index=['a', 'b']))
    return pd.DataFrame(means, index=group.index)

# 分组计算并按原顺序整理结果
result = df_with_idx.groupby('week').apply(calc_lagged_expanding_mean).reset_index(level=0, drop=True)
result.index = df.index  # 恢复原索引

print(result)

执行后输出：

a    b
week          
8     NaN  NaN
9     NaN  NaN
9     NaN  NaN
9     3.0  8.0
10    NaN  NaN
9     5.0  5.0
9     5.0  5.0

代码解释 #

1. 添加分组内行号：用cumcount()为每个week分组内的行从0开始编号，方便后续筛选行。
2. 自定义计算函数：对每个分组，遍历每行的行号，筛选出分组内行号小于当前行号-n的所有行，计算这些行的均值；如果没有符合条件的行，返回NaN。
3. 恢复原索引：计算完成后将结果的索引恢复为原DataFrame的索引，保证行顺序不变。

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