如何用Pandas按用户计算指定测试月的后续4个月登录滚动均值

需求说明 #

• 数据集包含用户月度登录次数，部分月份标记为test=True
• 需为每个test=True的月份，计算包含自身在内的后续4个月的登录数均值
• 核心约束：必须按用户独立统计，绝对不能跨用户取数；若用户数据末尾不足4个月，仅统计现有有效行

示例数据 #

import pandas as pd

# 示例数据集
df = pd.DataFrame({
    'user_id': ['A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B'],
    'month': ['2023-01', '2023-02', '2023-03', '2023-04', '2023-05', '2023-06', '2023-01', '2023-02', '2023-03', '2023-04'],
    'login_count': [10, 15, 20, 5, 8, 12, 7, 9, 14, 6],
    'test': [False, True, False, False, True, False, True, False, False, False]
})

常见错误实现（反面例子） #

# 错误：未按用户分组，导致跨用户统计；且未仅针对test=True行处理
df['test_avg'] = df[df['test']].rolling(window=4, min_periods=1)['login_count'].mean()

错误点：

• 没有按user_id分组，滚动窗口会跨用户取数
• 直接对筛选后的行计算，结果无法正确映射回原数据集的对应位置

期望输出 #

正确实现代码 #

# 1. 先按用户和月份排序（必须保证月份顺序正确，否则窗口计算无效）
df = df.sort_values(['user_id', 'month']).reset_index(drop=True)

# 2. 按用户分组计算滚动4期（含当前行）的均值，min_periods=1适配末尾不足4个月的情况
df['test_avg'] = df.groupby('user_id')['login_count'].transform(
    lambda x: x.rolling(window=4, min_periods=1).mean()
)

# 3. 仅保留test=True行的均值，其余行置为缺失值
df.loc[~df['test'], 'test_avg'] = pd.NA

代码关键说明 #

1. 排序：必须先按user_id和month排序，否则滚动窗口会取到乱序的月份数据，结果完全错误
2. 分组transform：groupby('user_id')确保每个用户的统计独立，transform将分组计算的结果映射回原DataFrame的对应行，避免索引错位
3. 滚动参数：window=4指定窗口包含当前行及后续3行（共4个月），min_periods=1允许窗口中至少有1行就计算均值，完美解决用户数据末尾不足4个月的场景
4. 过滤赋值：最后通过loc筛选出非测试行，将其test_avg设为缺失值，完全符合需求

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