示例

from sklearn.impute import IterativeImputer, KNNImputer, SimpleImputer

from datetime import datetime
def advanced_imputation(X_train, X_test, num_cols, method='iterative'):
    """
    高级缺失填充 - 只适用于数值列，必须在编码前执行！

    Parameters:
    -----------
    X_train, X_test : pd.DataFrame
        训练集和测试集
    num_cols : list
        需要处理的数值列名列表
    method : str
        填充方法: 'iterative', 'knn', 'simple'

    Returns:
    --------
    X_train, X_test : pd.DataFrame
        处理后的数据集（副本）
    """
    # 1. 创建副本，避免副作用
    X_train = X_train.copy()
    X_test = X_test.copy()

    # 2. 参数验证
    if not num_cols:
        print("  ⚠️ 未指定数值列")
        return X_train, X_test

    # 3. 过滤有效数值列
    valid_num_cols = []
    for col in num_cols:
        if col not in X_train.columns:
            print(f"  ⚠️ 列不存在: '{col}'")
            continue
        if not pd.api.types.is_numeric_dtype(X_train[col]):
            print(f"  ⚠️ 跳过非数值列: '{col}'")
            continue
        valid_num_cols.append(col)

    if not valid_num_cols:
        print("  ⚠️ 没有有效的数值列")
        return X_train, X_test

    # 4. 检查缺失值（训练集和测试集）
    has_missing_train = X_train[valid_num_cols].isnull().any().any()
    has_missing_test = X_test[valid_num_cols].isnull().any().any()

    if not has_missing_train and not has_missing_test:
        print(f"  ℹ️ 数值列无缺失值，跳过填充")
        return X_train, X_test

    # 5. 选择并执行填充方法
    try:
        if method == 'iterative':
            imputer = IterativeImputer(
                max_iter=10,
                random_state=42,
                n_nearest_features=None
            )
            print(f"  🔄 使用 IterativeImputer (迭代插补)")

        elif method == 'knn':
            from sklearn.impute import KNNImputer
            imputer = KNNImputer(
                n_neighbors=5,
                weights='uniform'
            )
            print(f"  🔄 使用 KNNImputer (K近邻插补)")
        else:
            raise ValueError(f"不支持的方法: {method}. 支持: 'iterative', 'knn'")

        # 6. 执行填充（使用 loc）
        X_train.loc[:, valid_num_cols] = imputer.fit_transform(
            X_train[valid_num_cols]
        )
        X_test.loc[:, valid_num_cols] = imputer.transform(
            X_test[valid_num_cols]
        )

        # 统计填充结果
        filled_train = X_train[valid_num_cols].isnull().sum().sum()
        filled_test = X_test[valid_num_cols].isnull().sum().sum()
        print(f"  ✅ 填充完成: {len(valid_num_cols)} 列")
        if filled_train > 0 or filled_test > 0:
            print(f"     ⚠️ 仍有缺失值: 训练集 {filled_train} 个, 测试集 {filled_test} 个")

    except Exception as e:
        print(f"  ❌ 填充失败: {e}")
        print(f"  ⚠️ 回退到中位数填充")
        imputer = SimpleImputer(strategy='median')
        X_train.loc[:, valid_num_cols] = imputer.fit_transform(
            X_train[valid_num_cols]
        )
        X_test.loc[:, valid_num_cols] = imputer.transform(
            X_test[valid_num_cols]
        )

    return X_train, X_test

调用示例

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

# 1. 创建示例数据
data = pd.DataFrame({
    'age': [25, 30, np.nan, 35, 40, 45],
    'income': [50000, 60000, 55000, np.nan, 70000, 65000],
    'score': [85, 90, 88, 92, np.nan, 95],
    'gender': ['M', 'F', 'M', 'F', 'M', 'F'],  # 分类变量，不会被填充
    'city': ['NY', 'LA', 'NY', 'LA', 'NY', 'LA']  # 分类变量，不会被填充
})

# 2. 划分训练集和测试集
X_train, X_test = train_test_split(data, test_size=0.3, random_state=42)

# 3. 指定数值列
num_cols = ['age', 'income', 'score']

# 4. 调用高级填充
X_train_imp, X_test_imp = advanced_imputation(
    X_train, X_test,
    num_cols=num_cols,
    method='knn'  # 使用迭代插补
)

print("填充前的训练集:")
print(X_train)
print("\n填充前的测试集:")
print(X_test)

print("填充后的训练集:")
print(X_train_imp)
print("\n填充后的测试集:")
print(X_test_imp)

IterativeImputer 和 KNNImputer 两种填充方法的核心区别

1. IterativeImputer（迭代插补）

python

# 原理：将每个缺失值列作为目标变量，用其他列作为特征进行预测
# 过程：
# Step 1: 用简单方法（如中位数）初步填充所有缺失值
# Step 2: 对每一列，用其他列作为特征，该列作为目标，训练模型
# Step 3: 用模型预测该列的缺失值
# Step 4: 重复 Step 2-3 多次（max_iter），直到收敛

可视化流程：

text

原始数据 → 初步填充（均值/中位数）
    ↓
循环迭代（max_iter次）：
    对每一列 j：
        用其他列 → 预测列j的缺失值
    ↓
检查收敛（变化是否小于 tol）
    ↓
返回最终结果

2. KNNImputer（K近邻插补）

python

# 原理：找到与缺失值样本最相似的K个邻居，用它们的平均值填充
# 过程：
# Step 1: 计算所有样本间的距离（考虑所有特征）
# Step 2: 对每个缺失值，找到K个最近邻
# Step 3: 用这K个邻居在该特征上的平均值填充

可视化流程：

text

对每个有缺失值的样本：
    计算与其他所有样本的距离（基于非缺失特征）
    ↓
    找到K个最近邻
    ↓
    用K个邻居在该特征的均值/中位数填充

详细对比表