熊猫：使用具有变量的组归纳描述性统计数据

Question

我有这样的数据框架：

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

我想做的是根据性别和班级列将年龄的缺失价值算。 我尝试使用一个函数，条件_impute进行操作。该功能的作用是采用数据框架和条件，然后使用它根据性别和班级分组将年龄归为年龄。但是警告是，条件可以是平均值或中位数，如果不是这两个中的任何一个，则该功能必须引起错误。 所以我这样做了：

### 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