Sklearn KNN Imputer按patient_id填充DataFrame时遇列数异常报错求助

1. KNNImputer仅处理数值型列
   KNNImputer只会对数值型特征执行插补操作，自动忽略字符串、类别型等非数值列（比如你的patient_id、visit_id可能属于这类）。原DataFrame的1198列中存在11列非数值列，因此fit_transform返回的数组仅包含1187个数值列，但你用原DataFrame的全部列名去创建新DataFrame，就会触发列数不匹配的报错。

2. 循环中merge操作未生效
   你的代码里data_imputed.merge(...)没有将结果赋值回data_imputed，导致每次循环后data_imputed始终是空DataFrame；同时merge的方式也不适合这种按患者拆分处理再合并的场景，拼接（concat）或分组处理更合理。

修正后的代码（分组处理更高效） #

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

# 初始化KNN插补器
knn = KNNImputer(missing_values=np.nan)

# 分离数值列（需要插补的医学分析结果）和非数值列
non_numeric_cols = data.select_dtypes(exclude=['number']).columns.tolist()
numeric_cols = data.select_dtypes(include=['number']).columns.tolist()

# 定义分组插补的函数
def impute_patient_group(group):
    # 对当前患者的数值列执行KNN插补
    imputed_numeric_data = knn.fit_transform(group[numeric_cols])
    # 将插补后的数值列替换回原分组
    group[numeric_cols] = pd.DataFrame(imputed_numeric_data, columns=numeric_cols, index=group.index)
    return group

# 按patient_id分组，应用插补函数
data_imputed = data.groupby('patient_id').apply(impute_patient_group)

代码说明 #

• 列类型分离：通过select_dtypes精准区分需要插补的数值列和无需处理的非数值列，避免KNNImputer跳过列导致的列数不匹配问题。
• 分组处理逻辑：用groupby('patient_id').apply替代手动循环，确保每个患者的样本仅在组内选取邻居进行插补，完全符合你的业务需求，同时代码更简洁高效。
• 保留原始结构：插补后直接替换组内的数值列，自动保留所有原始列、索引和非数值数据，无需额外合并操作。

额外提示 #

• 如果某个患者的样本数少于KNNImputer默认的n_neighbors=5，插补器会自动使用该患者的所有可用样本，不会报错；你也可以手动调整邻居数，比如KNNImputer(n_neighbors=3)。

内容的提问来源于stack exchange，提问作者fan-yang