如何在Pandas DataFrame中按国家计算6个月移动平均并补全NaN

现有以date为索引的Pandas DataFrame如下：

print(df)

               country        values           
date                        
2022-06-01     ES             2                  
2022-07-01     ES             2        
2022-08-01     ES             3         
2022-09-01     ES             3  
2022-10-01     ES             5  
2022-11-01     ES             5  
2022-12-01     ES             6  
2023-01-01     ES             6  
2023-02-01     ES             NaN 
2022-06-01     IT             1                    
2022-07-01     IT             1  
2022-08-01     IT             2  
2022-09-01     IT             2  
2022-10-01     IT             3  
2022-11-01     IT             3  
2022-12-01     IT             4 
2023-01-01     IT             4
2023-02-01     IT             NaN

需求 #

1. 生成名为moving_average的新列：按国家分组计算包含当前月份在内的过去6个月均值，不足6个月数据的位置填NaN；
2. 用对应月份的moving_average补全values列中的NaN值，且补全后的值会参与后续月份的移动平均计算。

预期输出如下：

country        values      moving_average 
date                        
2022-06-01     ES             2           NaN         
2022-07-01     ES             2           NaN
2022-08-01     ES             3           NaN
2022-09-01     ES             3           NaN
2022-10-01     ES             5           NaN
2022-11-01     ES             5           3.33  # AVG(2,2,3,3,5,5)
2022-12-01     ES             6           4      # AVG(2,3,3,5,5,6)
2023-01-01     ES             6           4.67   # AVG(3,3,5,5,6,6)，该值用于补全2023-02-01的NaN
2023-02-01     ES             4.67        4.945  # AVG(3,5,5,6,6,4.67)
2022-06-01     IT             1           NaN         
2022-07-01     IT             1           NaN
2022-08-01     IT             2           NaN
2022-09-01     IT             2           NaN
2022-10-01     IT             3           NaN
2022-11-01     IT             3           2      # AVG(1,1,2,2,3,3)
2022-12-01     IT             4           2.5    # AVG(1,2,2,3,3,4)
2023-01-01     IT             4           3      # AVG(2,2,3,3,4,4)，该值用于补全2023-02-01的NaN
2023-02-01     IT             3           3.167  # AVG(2,3,3,4,4,3)

这个需求的核心是迭代式补全NaN并计算移动平均，因为补全后的值会影响后续的均值计算，无法直接用普通的滚动窗口函数一步完成。我们可以按国家分组后，逐行处理每个分组的数据：

import pandas as pd

# 确保date索引为datetime类型（若原始数据不是则执行）
df.index = pd.to_datetime(df.index)

# 定义分组处理函数
def process_country_group(group):
    group = group.copy()
    group['moving_average'] = float('nan')
    
    # 逐行遍历计算并补全
    for idx in range(len(group)):
        # 截取当前行及前5行的窗口（共6行）
        window_data = group['values'].iloc[max(0, idx-5):idx+1]
        
        if len(window_data) == 6:
            current_avg = window_data.mean()
            # 保留三位小数对齐预期输出
            group['moving_average'].iloc[idx] = round(current_avg, 3)
            
            # 补全下一行的NaN值（如果存在）
            if idx + 1 < len(group) and pd.isna(group['values'].iloc[idx+1]):
                group['values'].iloc[idx+1] = current_avg
    
    return group

# 应用分组处理并合并结果
df = df.groupby('country', group_keys=False).apply(process_country_group)

代码说明 #

1. 索引类型校验：先将date索引转为datetime类型，避免窗口截取时出现逻辑错误；
2. 分组处理逻辑：对每个国家的子数据集单独处理：
   • 初始化moving_average列为NaN；
   • 遍历每一行，截取当前行及前5行的values数据作为计算窗口；
   • 当窗口长度刚好为6时，计算均值并赋值给当前行的moving_average；
   • 如果下一行的values是NaN，用当前计算的均值补全该值；
3. 结果合并：将处理后的各个国家分组合并为最终的DataFrame。

运行上述代码后，即可得到符合预期的输出结果。

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