如何用Pandas基于训练集Class分组均值填充测试集X列缺失值

步骤说明 #

1. 计算训练集分组均值：按Class分组，计算每组X列的均值（自动忽略分组内的NaN）
2. 映射均值填充测试集：将测试集的Class映射到对应均值，填充X列的缺失值

代码实现 #

首先导入依赖库并构建示例数据：

import pandas as pd
import numpy as np

# 训练集数据
train_df = pd.DataFrame({
    'Class': ['A', 'A', 'A', 'B', 'B', 'B'],
    'X': [10, np.nan, 20, 15, 17, np.nan]
})

# 测试集数据
test_df = pd.DataFrame({
    'Class': ['A', 'A', 'B', 'B'],
    'X': [11, np.nan, 25, np.nan]
})

计算训练集各分类的X均值：

# 按Class分组计算X列的均值（自动忽略NaN）
class_x_means = train_df.groupby('Class')['X'].mean()
# 输出结果：Class A为15，Class B为16
print(class_x_means)

填充测试集的缺失值：

# 用Class映射对应均值，填充X列的NaN
test_df['X'] = test_df['X'].fillna(test_df['Class'].map(class_x_means))

# 查看填充后的测试集
print(test_df)

最终效果 #

填充后的测试集如下：

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