Pandas DataFrame 按Subject分组结合条件填充height列缺失值的方法

核心逻辑梳理 #

完全匹配题目给定的规则，同时保证同一Subject对应的height值完全统一：

• 预先计算所有x == 'AA'样本的height列非空值中位数，作为2个空值分组的填充依据
• 按Subject分组统计组内的空值数量、组内已有的非空height取值
• 按规则匹配每个分组的统一height值：
  • 组内仅1个空值：直接使用组内已有的非空值填充全组
  • 组内存在2个空值：使用预计算的x == 'AA'组height中位数填充全组
• 将分组计算好的统一值映射回原表，完成列替换

可运行完整代码 #

import pandas as pd
import numpy as np

# 构造题目给定的原始数据
df = pd.DataFrame({
    'Subject': [1,1,2,2,3,3],
    'x': ['AA','AA','BB','BB','AA','AA'],
    'height': [130, np.nan, np.nan, 170, np.nan, np.nan]
})

# 预计算x=='AA'样本的height非空值中位数
aa_height_median = df.loc[df['x'] == 'AA', 'height'].median()

# 按Subject分组，统计组内空值数、组内有效非空值
group_stat = df.groupby('Subject')['height'].agg(
    null_cnt = lambda s: s.isna().sum(),
    valid_value = lambda s: s.dropna().iloc[0] if s.notna().any() else np.nan
).reset_index()

# 按规则生成每个Subject对应的统一height值
group_stat['final_height'] = np.where(
    group_stat['null_cnt'] == 2,
    aa_height_median,
    group_stat['valid_value']
)

# 映射回原表得到最终结果
result = df.drop(columns=['height']).merge(
    group_stat[['Subject', 'final_height']],
    on='Subject'
).rename(columns={'final_height': 'height'})

最终输出结果 #

运行代码后得到的result表内容：

结果校验：

• Subject1组共1个空值，使用组内非空值130填充，同组height值统一
• Subject2组共1个空值，使用组内非空值170填充，同组height值统一
• Subject3组共2个空值，使用x=='AA'组的height中位数130填充，同组height值统一
  完全符合题目要求的规则与约束。

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