Pandas：使用带有变量的 groupby 来估算描述性统计数据

Question

我有一个这样的数据框：

input_df = pd.DataFrame({"sex": ["M", "F", "F", "M", "M"], "Class": [1, 2, 2, 1, 1], "Age":[40, 30, 30, 50, NaN]})

我想做的是根据性别和 class 列估算年龄的缺失值。我已经尝试使用函数 conditional_impute 来完成它。该函数的作用是获取一个数据框和一个条件，然后使用它来根据性别和 class 分组估算年龄。但需要注意的是，条件可以是均值或中值，如果不是这两者中的任何一个，函数就必须引发错误。所以我这样做了：

### START FUNCTION
def conditional_impute(input_df, choice='median'):
    my_df = input_df.copy()
    # if choice is not median or mean, raise valueerror
    if choice == "mean" or choice == "median":
        my_df['Age'] = my_df['Age'].fillna(my_df.groupby(["Sex","Pclass"])['Age'].transform(choice))
    else:
        raise ValueError()    
    # round the values in Age colum
    my_df['Age'] = round(my_df['Age'], 1)
    return my_df
### END FUNCTION

但是我调用的时候出错了

conditional_impute(train_df, choice='mean')

我可能做错了什么？我实在搞不定。

Answer 1

如果您提供正确的输入，它的输出就很好...

# Fixed input to match function:
df = pd.DataFrame({"Sex": ["M", "F", "F", "M", "M"], "Pclass": [1, 2, 2, 1, 1], "Age":[40, 30, 30, 50, np.nan]})

def conditional_impute(input_df, choice='median'):
    my_df = input_df.copy()
    # if choice is not median or mean, raise valueerror
    if choice == "mean" or choice == "median":
        my_df['Age'] = my_df['Age'].fillna(my_df.groupby(["Sex","Pclass"])['Age'].transform(choice))
    else:
        raise ValueError()    
    # round the values in Age colum
    my_df['Age'] = round(my_df['Age'], 1)
    return my_df

conditional_impute(df, choice='mean')

输出：

  Sex  Pclass   Age
0   M       1  40.0
1   F       2  30.0
2   F       2  30.0
3   M       1  50.0
4   M       1  45.0

Pandas：使用带有变量的 groupby 来估算描述性统计数据

Pandas: imputing descriptive stats using a groupby with a variable

transform

pandas