使用Pandas的KNNImputer补全缺失数据时如何选取合适的K值

现有代码误差一致的原因 #

• 验证样本量太小：仅随机删除10个单元格，统计意义极弱，不同K值的效果差异会被随机误差完全覆盖，自然呈现误差一致的结果
• 误差计算逻辑错误：sum(abs(compare_df-transformed_df))默认按列求和，返回的是每列的误差总和而非10个缺失点的总误差，结果没有参考性
• 未固定随机种子：每次循环随机删除的位置不同，不同K值的验证集不统一，结果不具备可比性
• 验证逻辑有偏差：用KNN填充后的结果作为真值，相当于用模型输出验证模型本身，会引入额外偏差
• 未做特征标准化：KNN基于距离计算近邻，不同量纲的特征会导致距离计算完全失效，进一步抹平不同K值的效果差异

推荐的K值选取方法 #

方法1：带掩码的交叉验证（无后续建模任务时使用） #

核心逻辑是从原始数据中随机选择原本非空的单元格作为验证集，手动将其设为缺失值，用填充值和真实值的误差作为评估指标，所有K值在同一个验证集上测试保证可比性。
示例代码如下：

import numpy as np
import pandas as pd
from sklearn.impute import KNNImputer
from sklearn.preprocessing import StandardScaler

# 第一步：标准化特征（必须做，否则KNN效果完全不准）
scaler = StandardScaler()
X_scaled = scaler.fit_transform(ma_water_numeric)

# 第二步：构造统一的验证掩码
original_mask = ~np.isnan(X_scaled) # 原始数据中所有非空位置的掩码
total_valid = np.sum(original_mask)
n_validation = int(total_valid * 0.03) # 取3%的非空值作为验证集，可自行调整
# 固定随机种子保证可复现
rng = np.random.default_rng(seed=42)
# 随机选择验证位置
valid_flat_idx = rng.choice(np.where(original_mask.flatten())[0], size=n_validation, replace=False)
valid_rows, valid_cols = np.unravel_index(valid_flat_idx, X_scaled.shape)

# 第三步：构造带模拟缺失的数据集
X_masked = X_scaled.copy()
X_masked[valid_rows, valid_cols] = np.nan

# 第四步：遍历候选K值计算误差
candidate_k = [3,5,7,9,11,15,21,30]
error_list = []
for k in candidate_k:
    imputer = KNNImputer(n_neighbors=k)
    X_filled = imputer.fit_transform(X_masked)
    # 仅计算验证位置的平均绝对误差
    mae = np.mean(np.abs(X_filled[valid_rows, valid_cols] - X_scaled[valid_rows, valid_cols]))
    error_list.append((k, mae))

# 选取误差最小的K
best_k = sorted(error_list, key=lambda x:x[1])[0][0]

如果需要还原回原始量纲，只需要对填充后的数据集调用scaler.inverse_transform()即可。

方法2：绑定下游任务的网格搜索（有后续建模任务时使用） #

如果填充数据是为了后续的分类、回归等建模任务，可以将KNNImputer和下游模型绑定为Pipeline，用交叉验证同时搜索K值和模型参数，选出来的K最适配最终业务目标。
示例代码如下：

from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import GridSearchCV
from sklearn.preprocessing import StandardScaler

pipe = Pipeline([
    ('scaler', StandardScaler()),
    ('imputer', KNNImputer()),
    ('model', RandomForestRegressor(random_state=42))
])

param_grid = {
    'imputer__n_neighbors': [3,5,7,9,11,15],
    'model__n_estimators': [100,200]
}

# 5折交叉验证搜索最优参数，评估指标可根据任务替换
grid = GridSearchCV(pipe, param_grid, cv=5, scoring='neg_mean_absolute_error', n_jobs=-1)
grid.fit(X, y) # X为原始特征矩阵，y为下游任务的目标变量
print("最优K值：", grid.best_params_['imputer__n_neighbors'])

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