Pandas DataFrame - KNNImputer 算法实现

Question

我有一个包含缺失值的数据集，我想按组填充它们。我使用了 groupby() 方法，效果很好。但我想使用 KNNImputer 算法做同样的事情。

到目前为止我已经完成的代码（使用 groupby() 方法按预期工作）：

null_columns = df.columns[df.isnull().any()]

# filling median values by country
for column in null_columns:
    if column != "Life expectancy":
        df[column] = df.groupby("Country")[column].apply(lambda x: x.fillna(x.median()))

我尝试过但无法工作的代码（使用 KNNImputer）：

# Initializing imputer
imputer = KNNImputer(n_neighbors=5)

# Select numeric columns
numeric_cols = df.select_dtypes(include="number").columns

# Loop through numeric columns
for cols in numeric_cols:
    if cols != "Life expectancy":
        # Group by country and apply the imputer
        df[cols] = df.groupby("Country")[cols].transform(lambda x: imputer.fit_transform(x[[cols]]))

我尝试了不同的过程，但没有给我像前面的代码（使用 groupby() 方法）那样的结果。

Answer 1

KNNImputer 需要一个 2D 数组。您一次给出一列。
相反，将输入器应用于整个数据集。

循环每个国家/地区组，应用输入器，然后连接结果。
这是一个（未经测试的）示例：

def impute_groupwise(group):
    group_numeric = group[numeric_cols]
    imputed_values = imputer.fit_transform(group_numeric)
    group[numeric_cols] = imputed_values
    return group

df = df.groupby("Country").apply(impute_groupwise)