如何像groupby.agg()一样一步为Pandas添加行聚合列

问题描述 #

现有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({'ID':[1, 2, 3, 4, 5],
                   '1' :[7, 8, 2, 3, 0],
                   '2' :[1, 0, 4, 9, 9],
                   '3' :[3, 1, 1, 6, 8]})

输出结果：

ID  1  2  3
0   1  7  1  3
1   2  8  0  1
2   3  2  4  1
3   4  3  9  6
4   5  0  9  8

需要为每行添加基于列'1'、'2'、'3'的Av（均值）、Sum（求和）、Count（计数）、Max（最大值）列，当前实现方式如下：

cols = ['1', '2', '3']
df['Av'] = df[cols].mean(axis=1)
df['Sum'] = df[cols].sum(axis=1).fillna(0)
df['Count'] = df[cols].count(axis=1).fillna(0)
df['Max'] = df[cols].max(axis=1).fillna(0)

输出结果：

ID  1  2  3        Av  Sum  Count  Max
0   1  7  1  3  3.666667   11      3    7
1   2  8  0  1  3.000000    9      3    8
2   3  2  4  1  2.333333    7      3    4
3   4  3  9  6  6.000000   18      3    9
4   5  0  9  8  5.666667   17      3    9

请问是否有更简洁的方式，比如像groupby.agg()那样用1-2行代码完成上述操作？

简洁实现方案 #

当然可以，利用DataFrame.agg()并指定axis=1，就能一次性计算所有需要的行级统计量，再合并到原DataFrame中，一行代码即可完成：

方法一：直接覆盖原DataFrame #

cols = ['1', '2', '3']
df = df.join(df[cols].agg(['mean', 'sum', 'count', 'max'], axis=1).rename(columns={'mean':'Av', 'sum':'Sum', 'count':'Count', 'max':'Max'}))

方法二：生成新DataFrame（保留原数据） #

如果不想修改原DataFrame，用assign更合适：

cols = ['1', '2', '3']
df_new = df.assign(**df[cols].agg(['mean', 'sum', 'count', 'max'], axis=1).rename(columns={'mean':'Av', 'sum':'Sum', 'count':'Count', 'max':'Max'}))

处理缺失值的情况 #

如果你的实际数据存在缺失值，只需在agg后链式调用.fillna(0)即可：

df = df.join(df[cols].agg(['mean', 'sum', 'count', 'max'], axis=1)
             .rename(columns={'mean':'Av', 'sum':'Sum', 'count':'Count', 'max':'Max'})
             .fillna(0))

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