如何修改PLUMBER类实现pandas DataFrame指定列按分组均值替换NaN

代码触发报错的核心原因及其他逻辑问题如下：

• groupby参数传值错误：label_column_names本身就是存储分组列名的列表，额外套一层[]会变成嵌套列表，导致pandas无法识别分组键，直接触发长度不匹配报错。
• 类初始化逻辑错误：__init__方法没有设置入参，直接引用外部变量赋值属性，不符合封装逻辑，运行时会直接报变量未定义错误。
• fill_na方法逻辑冗余：方法开头强制把传入的参数全部替换为实例属性，传入的参数完全失效，逻辑矛盾。
• 自定义均值方法无边界处理：如果某分组下指定数值列全为NaN，会触发除零报错。
• 列名处理未生效：原代码对列名做了去空格处理，但没有重新赋值给df的列属性，会导致带空格的列名匹配失败。

import pandas as pd
import numpy as np

class PLUMBER():
    def __init__(self, df, numerical_feature_names, label_column_names):
        # 初始化时接收外部传入的参数赋值给实例属性
        self.df = df
        self.numerical_feature_names = numerical_feature_names
        self.label_column_names = label_column_names

    def mean(self, nums):
        # 过滤掉空值，避免除零错误
        valid_nums = [num for num in nums if pd.notna(num)]
        if not valid_nums:
            return pd.NA
        return sum(valid_nums) / len(valid_nums)

    def fill_na(self):
        # 直接使用实例属性，无需重复传参
        df = self.df.copy() # 加copy避免修改原df，不需要可直接删掉该行
        for numerical_feature_name in self.numerical_feature_names:
            # 去掉groupby参数外层多余的方括号
            df[numerical_feature_name] = df.groupby(self.label_column_names)[numerical_feature_name].transform(lambda x: x.fillna(self.mean(x)))
        return df

if __name__=="__main__":
    # 测试数据
    d={'month': ['01/01/2020', '01/02/2020', '01/03/2020', '01/01/2020', '01/02/2020', '01/03/2020'], 
       'country': ['Japan', 'Japan', 'Japan', 'Poland', 'Poland', 'Poland'], 
       'level':['A01', 'A01', 'A01', 'A00','A00', 'A00'],
       'job title':['Insights Manager', 'Insights Manager', 'Insights Manager', 'Sales Director', 'Sales Director', 'Sales Director'],
       'number':[np.nan, 450, 299, np.nan, 19, 29],
       'age':[np.nan, 30, 28, np.nan, 29, 18]}
    df=pd.DataFrame(d)

    # 列名处理
    column_names=df.columns.values.tolist()
    column_names= [column_name.strip() for column_name in column_names]
    df.columns = column_names # 给df重新赋值处理后的列名，避免列名匹配失败

    # 参数定义
    label_column_names=['country', 'level', 'job title']
    numerical_feature_names = [x for x in column_names if x not in label_column_names]
    numerical_feature_names.remove('month')

    # 初始化类传入参数
    plumber=PLUMBER(df, numerical_feature_names, label_column_names)
    # 执行空值填充
    df_filled = plumber.fill_na()
    print(df_filled)

修改后运行得到的df中：

• 日本组的number空值替换为374.5，age空值替换为29
• 波兰组的number空值替换为24，age空值替换为23.5
  完全符合按分组均值替换空值的需求。

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