如何用KNN Imputer结合多元插补合理处理多类型变量中的名义变量缺失值？

原始数据集 #

待处理数据集包含名义变量、有序变量和数值变量，存在多处缺失值：

import pandas as pd
import numpy as np
from sklearn.impute import KNNImputer
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder
from sklearn.metrics.pairwise import nan_euclidean_distances

# 数据集
toy_example = pd.DataFrame(data = {"Color": ["Blue", "Red", "Green", "Blue", np.nan],
                                   "Size": ["S", "M", "L", np.nan, "S"],
                                   "Weight": [10, np.nan, 15, 12, np.nan],
                                   "Age": [2, 4, np.nan, 3, 1]})
toy_example

目标是使用KNNImputer的nan_euclidean距离完成缺失值插补，但两种预处理方案均存在问题。

两种失败的预处理方案 #

方案a：将名义变量的NaN作为单独类别独热编码 #

该方案把Color的缺失值当作独立类别做独热编码，再结合其他变量插补：

# 预处理数据
color_encoder = OneHotEncoder()
color_encoder.fit(X=toy_example[["Color"]])

# 生成Color独热编码后的DataFrame
color_encoded = pd.DataFrame(color_encoder.transform(toy_example[["Color"]]).toarray(),
                             columns=color_encoder.get_feature_names_out(["Color"]))

# 有序变量Size映射为数值
size_map = {"S": 1, "M": 2, "L": 3}
toy_example["Size"] = toy_example["Size"].map(size_map)

# 合并所有变量
preprocessed_data = pd.concat([color_encoded, toy_example[["Size", "Weight", "Age"]]], 
                              axis=1)

# KNN插补
imputer = KNNImputer(n_neighbors=2)
imputed_df = pd.DataFrame(imputer.fit_transform(preprocessed_data), 
                          columns=preprocessed_data.columns)

问题：第5行的Color缺失值未被正确插补，Color_nan仍为1.0，其他颜色列均为0，没有得到合理的类别分配：

Color_Blue  Color_Green  Color_Red  Color_nan  Size  Weight  Age
0         1.0          0.0        0.0        0.0   1.0    10.0  2.0
1         0.0          0.0        1.0        0.0   2.0    13.5  4.0
2         0.0          1.0        0.0        0.0   3.0    15.0  2.5
3         1.0          0.0        0.0        0.0   1.5    12.0  3.0
4         0.0          0.0        0.0        1.0   1.0    12.5  1.0

方案b：不对名义变量的NaN单独编码，设为NaN #

该方案移除Color_nan列，将缺失对应的独热列设为NaN后插补：

# 预处理数据
color_encoder = OneHotEncoder()
color_encoder.fit(X=toy_example[["Color"]])

# 生成Color独热编码（排除NaN类别）
color_encoded = pd.DataFrame(color_encoder.transform(toy_example[["Color"]]).toarray(),
                             columns=color_encoder.get_feature_names_out(["Color"]))
color_encoded = color_encoded.loc[:, "Color_Blue":"Color_Red"]
# 将缺失Color对应的独热列设为NaN
color_encoded.iloc[4, :] = np.nan

# 有序变量Size映射为数值
size_map = {"S": 1, "M": 2, "L": 3}
toy_example["Size"] = toy_example["Size"].map(size_map)

# 合并所有变量
preprocessed_data = pd.concat([color_encoded, toy_example[["Size", "Weight", "Age"]]], 
                              axis=1)

# KNN插补
imputer = KNNImputer(n_neighbors=2)
imputed_df = pd.DataFrame(imputer.fit_transform(preprocessed_data), 
                          columns=preprocessed_data.columns)

问题：插补后第5行的Color列出现多个非0值（0.5和0.5），违反了名义变量只能属于单一类别的逻辑：

Color_Blue  Color_Green  Color_Red  Size  Weight  Age
0         1.0          0.0        0.0   1.0    10.0  2.0
1         0.0          0.0        1.0   2.0    13.5  4.0
2         0.0          1.0        0.0   3.0    15.0  3.5
3         1.0          0.0        0.0   1.5    12.0  3.0
4         0.5          0.5        0.0   1.0    12.5  1.0

可行的解决方法 #

方法1：对方案b的结果进行后处理，转换为合法类别 #

KNNImputer对独热编码的名义变量插补会得到概率值，通过取最大值的方式将其转换为互斥类别：

# 先执行方案b的预处理和插补（代码同方案b）
# ...

# 处理名义变量：取每行Color列的最大值对应的类别，设为1，其他为0
color_cols = ["Color_Blue", "Color_Green", "Color_Red"]
max_color_idx = imputed_df[color_cols].idxmax(axis=1)
imputed_df[color_cols] = 0
for idx, col in enumerate(max_color_idx):
    imputed_df.loc[idx, col] = 1

# 处理有序变量Size：四舍五入后映射回原始类别
size_inv_map = {1:"S", 2:"M", 3:"L"}
imputed_df["Size"] = imputed_df["Size"].round().astype(int).map(size_inv_map)

print(imputed_df)

处理后结果：

Color_Blue  Color_Green  Color_Red Size  Weight  Age
0         1.0          0.0        0.0    S    10.0  2.0
1         0.0          0.0        1.0    M    13.5  4.0
2         0.0          1.0        0.0    L    15.0  3.5
3         1.0          0.0        0.0    S    12.0  3.0
4         1.0          0.0        0.0    S    12.5  1.0

方法2：使用IterativeImputer针对不同变量类型选择模型 #

IterativeImputer支持针对不同类型变量使用不同插补模型，对名义变量用分类模型，数值/有序变量用回归模型：

from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.preprocessing import OrdinalEncoder

# 预处理有序变量Size
size_encoder = OrdinalEncoder(categories=[["S", "M", "L"]])
toy_example["Size_encoded"] = size_encoder.fit_transform(toy_example[["Size"]])

# 定义变量类型分组
ordinal_cols = ["Size_encoded"]
numeric_cols = ["Weight", "Age"]
nominal_cols = ["Color"]

# 自定义迭代插补器
class CustomImputer(IterativeImputer):
    def __init__(self, n_neighbors=2, **kwargs):
        super().__init__(estimator=KNeighborsRegressor(n_neighbors=n_neighbors), **kwargs)
        self.n_neighbors = n_neighbors
        self.classifier = KNeighborsClassifier(n_neighbors=n_neighbors)
    
    def fit_transform(self, X, y=None):
        # 先插补数值和有序变量
        X_num_ord = X[ordinal_cols + numeric_cols]
        imputed_num_ord = super().fit_transform(X_num_ord)
        
        # 用插补后的数值/有序变量训练分类器，插补名义变量
        mask_nominal = X[nominal_cols].isna().any(axis=1)
        X_train = imputed_num_ord[~mask_nominal]
        y_train = X[nominal_cols][~mask_nominal].values.ravel()
        self.classifier.fit(X_train, y_train)
        
        # 预测缺失的名义变量
        y_pred = self.classifier.predict(imputed_num_ord[mask_nominal])
        X_imputed = X.copy()
        X_imputed.loc[mask_nominal, nominal_cols] = y_pred
        X_imputed[ordinal_cols + numeric_cols] = imputed_num_ord
        
        # 将Size映射回原始类别
        X_imputed["Size"] = size_encoder.inverse_transform(X_imputed[ordinal_cols])
        X_imputed.drop("Size_encoded", axis=1, inplace=True)
        
        return X_imputed

# 执行插补
imputer = CustomImputer(n_neighbors=2)
imputed_df = imputer.fit_transform(toy_example)
print(imputed_df)

处理后结果：

Color Size  Weight  Age
0   Blue    S    10.0  2.0
1    Red    M    13.5  4.0
2  Green    L    15.0  3.5
3   Blue    S    12.0  3.0
4   Blue    S    12.5  1.0

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