Sklearn技巧：如何让OneHotEncoder保留NaN值并实现填充？

测试数据 #

import pandas as pd
import numpy as np
from sklearn.preprocessing import OneHotEncoder
from sklearn.impute import SimpleImputer

test_df = pd.DataFrame({'sex': ['female', 'male', np.nan], 'eye_color': ['blue', np.nan, 'brown']})

测试数据输出：

当前问题 #

使用默认配置的OneHotEncoder时，原始数据中的NaN会被全部转为0，无法保留缺失信息：

ohe = OneHotEncoder(
    categories=[['female', 'male'], ['blue', 'brown', 'green']],
    handle_unknown="ignore",
    sparse_output=False,
).set_output(transform='pandas')
current_df = ohe.fit_transform(test_df)

输出结果：

解决方案 #

核心思路是先记录原始数据中每个特征的NaN位置，编码完成后将对应特征的所有编码列恢复为NaN，具体实现如下：

方法1：基于sklearn OneHotEncoder的手动恢复 #

# 1. 执行编码
ohe = OneHotEncoder(
    categories=[['female', 'male'], ['blue', 'brown', 'green']],
    handle_unknown="ignore",
    sparse_output=False,
).set_output(transform='pandas')
encoded_df = ohe.fit_transform(test_df)

# 2. 建立原始特征与编码列的映射关系
feature_cols = ohe.feature_names_in_
encoded_col_map = {
    feat: [col for col in encoded_df.columns if col.startswith(f"{feat}_")]
    for feat in feature_cols
}

# 3. 根据原始NaN掩码恢复缺失值
for feat, cols in encoded_col_map.items():
    nan_mask = test_df[feat].isna()
    encoded_df.loc[nan_mask, cols] = np.nan

# 此时encoded_df即为保留NaN的目标结果
desired_df = encoded_df

输出结果：

方法2：使用pandas.get_dummies（更简洁） #

pandas.get_dummies默认会保留原始NaN值，只需补充缺失的类别列并调整格式即可：

# 生成编码列，指定预设类别
desired_df = pd.get_dummies(
    test_df,
    columns=['sex', 'eye_color'],
    categories={'sex': ['female', 'male'], 'eye_color': ['blue', 'brown']}
)

# 手动添加未出现在原始数据中的类别列
desired_df['eye_color_green'] = 0.0

# 调整列顺序与目标一致
desired_df = desired_df[['sex_female', 'sex_male', 'eye_color_blue', 'eye_color_brown', 'eye_color_green']]

# 恢复原始NaN位置的缺失值
desired_df.loc[test_df['sex'].isna(), ['sex_female', 'sex_male']] = np.nan
desired_df.loc[test_df['eye_color'].isna(), ['eye_color_blue', 'eye_color_brown', 'eye_color_green']] = np.nan

后续填充 #

得到保留NaN的编码结果后，即可用SimpleImputer完成填充：

ultimately_df = SimpleImputer().set_output(transform="pandas").fit_transform(desired_df)

最终填充结果：

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