Pandas按A1、A2分组后仅对B列非NaN值计算窗口3的滚动均值如何实现

实现方案 #

核心逻辑：滚动均值计算完成后，基于原始数据的空值标记，将原本就为NaN的位置手动覆盖为空即可，完整可运行代码如下：

import pandas as pd
import numpy as np

# 构造测试DataFrame
df = pd.DataFrame({'B': [0, 1, 2, np.nan, 4],
                   'A1': [1, 1, 2, 2, 2],
                   'A2': [1, 2, 3, 3, 3]})

# 按A1、A2分组后计算窗口为3的滚动均值，min_periods=1允许不足3个有效值时计算
roll_result = df.groupby(['A1', 'A2'])['B'].transform(lambda x: x.rolling(window=3, min_periods=1).mean())
# 原始B列为NaN的位置强制保留NaN，其余位置用滚动计算结果
df['B_roll_mean'] = roll_result.where(df['B'].notna(), np.nan)

运行后df['B_roll_mean']的输出结果和期望完全一致：[0.0, 1.0, 2.0, NaN, 3.0]

原方案问题说明 #

pandas的rolling.mean()在计算时默认会跳过窗口内的NaN值，不会保留原始列中的空值标记，因此直接计算会把原本是空值的位置填充为窗口内其余有效值的均值，额外加一层空值覆盖即可解决该问题。

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