使用KNNImputer填充缺失值时如何保留Pandas DataFrame列名？

嘿，这个问题我之前用KNNImputer的时候也踩过坑！核心原因是KNNImputer.fit_transform()返回的是NumPy数组，而不是Pandas DataFrame，所以原始列名自然就丢失了。至于你遇到的ValueError: could not convert string to float: 'age'，应该是你在尝试排除列的时候操作有误——不小心把列名的字符串列表直接传给了fit_transform()，但KNNImputer只能处理数值型的矩阵/数组，所以它会把'age'当成待转换的数据，自然报错。

下面给你两种清晰的解决方案，优先推荐第一种（简单直接，无需额外依赖）：

方法1：先存储原始列名，转换后重新赋值 #

这是最常用的办法，步骤超简单：

1. 提前保存原始DataFrame的列名；
2. 用KNNImputer处理数据得到数组；
3. 将数组转回DataFrame时，指定columns参数为原始列名。

完整代码示例：

import pandas as pd
from sklearn.impute import KNNImputer

# 你的原始数值DataFrame
numeric_data = pd.DataFrame({
    'age': [19, 18, 28, 33],
    'bmi': [None, 33.77, 33.0, 22.705],
    'children': [0.0, 1.0, 3.0, 0.0],
    'charges': [16884.9240, None, 4449.4620, None]
})

# 1. 保存原始列名
original_cols = numeric_data.columns

# 2. 初始化KNNImputer并执行填充
imputer = KNNImputer(n_neighbors=2, weights="uniform")
imputed_array = imputer.fit_transform(numeric_data)

# 3. 转回DataFrame并恢复列名
impute_data = pd.DataFrame(imputed_array, columns=original_cols)
print(impute_data)

运行后你会看到列名完美保留，和原始DataFrame一致。

方法2：使用Pipeline+ColumnTransformer（适合复杂场景） #

如果你的数据里混合了数值列和非数值列，需要只对数值列做填充，同时保留所有列的结构，可以用ColumnTransformer配合Pipeline，这样能全程保持DataFrame的结构（包括列名）。

示例代码：

import pandas as pd
from sklearn.impute import KNNImputer
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline

# 假设你的数据有非数值列（比如'gender'），这里模拟一下
data = pd.DataFrame({
    'gender': ['female', 'male', 'male', 'female'],
    'age': [19, 18, 28, 33],
    'bmi': [None, 33.77, 33.0, 22.705],
    'children': [0.0, 1.0, 3.0, 0.0],
    'charges': [16884.9240, None, 4449.4620, None]
})

# 定义要填充的数值列
numeric_features = ['age', 'bmi', 'children', 'charges']

# 创建填充器：只对数值列做KNN填充
numeric_transformer = Pipeline(steps=[
    ('imputer', KNNImputer(n_neighbors=2, weights="uniform"))
])

# 构建ColumnTransformer，指定数值列的处理方式，其他列保持原样
preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, numeric_features)
    ],
    remainder='passthrough'  # 保留非数值列
)

# 执行填充，结果转回DataFrame
imputed_data = pd.DataFrame(preprocessor.fit_transform(data), 
                            columns=numeric_features + ['gender'])
print(imputed_data)

这种方法适合数据预处理流程比较复杂的场景，能避免手动管理列名的麻烦。

最后再提一下你遇到的那个报错：只要确保传给fit_transform()的是数值型的DataFrame或数组，而不是列名字符串，就不会再出现could not convert string to float的问题啦。

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