使用Pandas按key_id分组，用近3个月均值填充缺失月份数据

现有一个包含key_id列、value列和date列的DataFrame，部分key_id分组下存在月份缺失的行。需为每个key_id的缺失月份，用其过去3个月的value均值填充对应行的value值。

示例输入数据 #

key_id  value  date
1      50.0    2021-01-01
1      54.0    2021-02-01
1      36.0    2021-04-01
1      48.5    2021-05-01
2      50.0    2021-01-01
2      54.0    2021-02-01
2      36.0    2021-03-01
2      48.5    2021-05-01

注：上述数据中，key_id=1缺失了2021-03-01的行，key_id=2缺失了2021-04-01的行，需用过去3个月的均值填充这些缺失行的value。

预期输出 #

key_id  value  date
1      50.0    2021-01-01
1      54.0    2021-02-01
1      52.0    2021-03-01  # 注：按实际过去2个月均值计算为52.0，示例中33.0应为笔误
1      36.0    2021-04-01
1      48.5    2021-05-01
2      50.0    2021-01-01
2      54.0    2021-02-01
2      36.0    2021-03-01
2      46.67   2021-04-01  # 注：按过去3个月均值计算约为46.67，示例中40.0应为笔误
2      48.5    2021-05-01

解决方案 #

以下是基于Pandas的实现代码，步骤清晰且符合需求：

import pandas as pd

# 构造示例DataFrame
data = {
    'key_id': [1, 1, 1, 1, 2, 2, 2, 2],
    'value': [50.0, 54.0, 36.0, 48.5, 50.0, 54.0, 36.0, 48.5],
    'date': ['2021-01-01', '2021-02-01', '2021-04-01', '2021-05-01',
             '2021-01-01', '2021-02-01', '2021-03-01', '2021-05-01']
}
df = pd.DataFrame(data)

# 统一日期格式为每月第一天，确保按月份对齐
df['date'] = pd.to_datetime(df['date']).dt.to_period('M').dt.to_timestamp()

# 定义分组处理函数：补全缺失月份并填充均值
def process_group(group):
    # 生成该组完整的月份序列（从最早到最晚日期的每个月第一天）
    full_dates = pd.date_range(start=group['date'].min(), end=group['date'].max(), freq='MS')
    # 重新索引，补全缺失行
    group = group.set_index('date').reindex(full_dates).reset_index().rename(columns={'index': 'date'})
    # 填充key_id（缺失行继承分组的key_id）
    group['key_id'] = group['key_id'].ffill()
    # 用过去最多3个月的均值填充缺失value，closed='left'确保只取当前日期之前的数据
    group['value'] = group['value'].rolling(window=3, min_periods=1, closed='left').mean()
    return group

# 对每个key_id分组应用处理函数
result_df = df.groupby('key_id', group_keys=False).apply(process_group)

# 打印结果（保留两位小数）
print(result_df.round(2))

代码说明 #

1. 日期格式统一：将date转换为每月第一天的格式，避免日期格式不一致导致的补全错误
2. 补全缺失月份：为每个key_id分组生成完整的月份序列，确保没有月份遗漏
3. 滚动均值填充：使用rolling窗口计算过去最多3个月的均值，closed='left'保证只使用当前日期之前的数据，符合“过去3个月”的要求；min_periods=1处理数据不足3个月的边界情况

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