如何使用SimpleImputer按行业分组均值填补员工数、预算列的NaN缺失值

结论 #

原生SimpleImputer无法直接实现按行业分组填充的需求，它仅支持全局维度的统计值填充，没有内置分组计算的逻辑。你可以用以下两种方案实现需求：

方案1：直接用Pandas实现（最简便） #

你当前把DataFrame转成了numpy数组X反而丢失了列名和分组关联信息，建议先在DataFrame层面处理完缺失值再转成数组：

import pandas as pd
import numpy as np

dataset = pd.read_csv('Group A.csv')
# 按行业分组，对员工数、预算两列用对应分组的均值填充缺失值
dataset[['# of Employees', '# Budget']] = dataset.groupby('Industry')[['# of Employees', '# Budget']].transform(
    lambda x: x.fillna(x.mean())
)
# 处理完缺失值后再生成你需要的特征数组X
X = dataset.iloc[:, :-1].values

如果存在某行业下某列所有值都是空的情况，可以加兜底逻辑，fallback到全局均值填充：

cols_to_impute = ['# of Employees', '# Budget']
global_means = dataset[cols_to_impute].mean()
dataset[cols_to_impute] = dataset.groupby('Industry')[cols_to_impute].transform(
    lambda x: x.fillna(x.mean() if not pd.isna(x.mean()) else global_means[x.name])
)

方案2：兼容Sklearn接口实现（需放入Pipeline流程时使用） #

如果你需要把填充逻辑嵌入Sklearn的Pipeline工作流，可以自定义分组填充转换器：

from sklearn.base import BaseEstimator, TransformerMixin
import pandas as pd

class GroupMeanImputer(BaseEstimator, TransformerMixin):
    def __init__(self, group_col, impute_cols):
        self.group_col = group_col
        self.impute_cols = impute_cols
        self.group_means_ = None
        self.global_means_ = None
    
    def fit(self, X, y=None):
        self.group_means_ = X.groupby(self.group_col)[self.impute_cols].mean()
        self.global_means_ = X[self.impute_cols].mean()
        return self
    
    def transform(self, X):
        X_copy = X.copy()
        for col in self.impute_cols:
            X_copy[col] = X_copy.apply(
                lambda row: self.group_means_.loc[row[self.group_col], col] 
                if pd.isna(row[col]) and not pd.isna(self.group_means_.loc[row[self.group_col], col])
                else self.global_means_[col] if pd.isna(row[col])
                else row[col], 
                axis=1
            )
        return X_copy

# 调用示例
dataset = pd.read_csv('Group A.csv')
imputer = GroupMeanImputer(group_col='Industry', impute_cols=['# of Employees', '# Budget'])
dataset = imputer.fit_transform(dataset)
X = dataset.iloc[:, :-1].values

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