如何在Pandas GroupBy中优雅合并Agg与Transform操作

现有如下DataFrame：

df = pd.DataFrame(
    {'stud_id' : [101, 101, 101, 101, 
                  101, 101, 101, 101],
     'sub_code' : ['CSE01', 'CSE01', 'CSE01', 
                   'CSE01', 'CSE02', 'CSE02',
                   'CSE02', 'CSE02'],
     'ques_date' : ['13/11/2020', '10/1/2018','11/11/2017', '27/03/2016', 
                '13/05/2010',  '10/11/2008','11/1/2007', '27/02/2006'],
     'marks' : [77, 86, 55, 90, 
                65, 90, 80, 67]}
)
df['ques_date'] = pd.to_datetime(df['ques_date'])

需要完成这些操作：

• 按stud_id和sub_code分组；
• 计算每组ques_date的日期差中位数；
• 计算每组marks的计数，同时获取每组最晚的ques_date。

目前用两行代码分别实现transform和agg操作：

df['avg_ques_gap'] = (df.groupby(['stud_id','sub_code'])['ques_date']
                          .transform(lambda x: x.diff().dt.days.median()))

output = df.groupby(['stud_id','sub_code']).agg(last_ques_date=('ques_date','max'),
                                                  total_pos_transactions=('marks','count')).reset_index()

想知道能不能把这两个操作合并成单行代码实现。

当然可以合并，这里提供两种单行实现的思路：

思路1：聚合后关联回原数据 #

直接在groupby.agg里同时计算所有需要的聚合值，再通过merge关联到原DataFrame，一行搞定：

result = df.merge(df.groupby(['stud_id', 'sub_code']).agg(last_ques_date=('ques_date', 'max'), total_pos_transactions=('marks', 'count'), avg_ques_gap=('ques_date', lambda x: x.diff().dt.days.median())).reset_index(), on=['stud_id', 'sub_code'])

思路2：链式调用生成汇总表 #

如果只需要最终的汇总结果（不需要保留原数据的每一行），可以用链式调用把transform和agg串起来：

output = df.assign(avg_ques_gap=df.groupby(['stud_id','sub_code'])['ques_date'].transform(lambda x: x.diff().dt.days.median())).groupby(['stud_id','sub_code']).agg(last_ques_date=('ques_date','max'), total_pos_transactions=('marks','count'), avg_ques_gap=('avg_ques_gap','first')).reset_index()

说明 #

• 思路1的结果会保留原DataFrame的所有行，同时添加三个聚合字段；
• 思路2直接生成你需要的汇总表，因为avg_ques_gap在每组里的值都相同，所以取first就能拿到正确的中位数。

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