Pandas按组填充缺失值：基于前后各2个周期有效数据的均值填充及标记

需求说明 #

现有分客户、分周期的业务数据，部分周期的volume字段存在缺失。需要：

• 仅当缺失值的前2个周期和后2个周期的volume数据均无缺失时，用这4个值的均值填充该缺失值
• 新增imputed字段：填充值标记为1，原始值（包括未满足条件的缺失值）标记为0

初始数据 #

import pandas as pd
df = pd.DataFrame({
    'cust_id': [1,1,1,1,1,1,2,2,2,2,2,2],
    'period' : [1,2,3,4,5,6,1,2,3,4,5,6],
    'volume' : [1,2,None,4,5,6,7,None,9,None,11,None]
})

规则匹配逻辑 #

• 客户1的第3周期缺失：前2个周期（1、2）和后2个周期（4、5）的volume都存在，符合填充条件，计算均值为(1+2+4+5)/4 = 3
• 客户2的3个缺失值：
  • 第2周期缺失：前2个周期仅第1周期存在，不符合
  • 第4周期缺失：后2个周期仅第5周期存在，不符合
  • 第6周期缺失：无后2个周期数据，不符合
    因此均不填充

实现代码 #

import pandas as pd

# 加载初始数据
df = pd.DataFrame({
    'cust_id': [1,1,1,1,1,1,2,2,2,2,2,2],
    'period' : [1,2,3,4,5,6,1,2,3,4,5,6],
    'volume' : [1,2,None,4,5,6,7,None,9,None,11,None]
})

# 初始化标记字段
df['imputed'] = 0

# 按客户ID分组处理每个客户的周期数据
def process_customer(group):
    # 遍历当前客户的所有缺失值行
    for idx, row in group[group['volume'].isna()].iterrows():
        current_p = row['period']
        # 筛选前后各2个周期的数据
        prev_rows = group[group['period'].isin([current_p-2, current_p-1])]
        next_rows = group[group['period'].isin([current_p+1, current_p+2])]
        
        # 检查前后各2个周期是否都有完整数据
        if (len(prev_rows) == 2 and prev_rows['volume'].notna().all() and
            len(next_rows) == 2 and next_rows['volume'].notna().all()):
            # 计算均值并填充
            avg_val = pd.concat([prev_rows['volume'], next_rows['volume']]).mean()
            group.loc[idx, 'volume'] = avg_val
            # 标记为填充值
            group.loc[idx, 'imputed'] = 1
    return group

# 应用分组处理并重置索引
df = df.groupby('cust_id').apply(process_customer).reset_index(drop=True)

# 输出结果
print(df)

运行结果 #

cust_id  period  volume  imputed
0         1       1     1.0        0
1         1       2     2.0        0
2         1       3     3.0        1
3         1       4     4.0        0
4         1       5     5.0        0
5         1       6     6.0        0
6         2       1     7.0        0
7         2       2     NaN        0
8         2       3     9.0        0
9         2       4     NaN        0
10        2       5    11.0        0
11        2       6     NaN        0

注：原问题示例中的期望结果将客户1第3周期填充为6是计算错误，正确均值应为3

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