如何从Pandas groupby聚合结果中提取指定列最大值对应的行

• 你调用agg(['count'])后得到的是多级列索引的DataFrame，列名格式为(原列名, 'count')，比如foo列的计数字段是('foo', 'count')
• 你可以直接基于这个多级索引筛选，也可以先把列名拍平简化后续操作

方案1：直接基于多级索引筛选 #

# 先保存分组统计结果
grouped_df = df.groupby('country').agg(['count'])

# 分别计算foo、bar列count的最大值
max_foo = grouped_df[('foo', 'count')].max()
max_bar = grouped_df[('bar', 'count')].max()

# 筛选满足 foo count等于最大值 或 bar count等于最大值的行
target_rows = grouped_df[(grouped_df[('foo', 'count')] == max_foo) | (grouped_df[('bar', 'count')] == max_bar)]

运行后target_rows就会返回USA、India对应的两行结果。

方案2：聚合时拍平列名（更易维护） #

如果你不想处理多级索引，可以在聚合阶段直接指定扁平化的列名：

grouped_df = df.groupby('country', as_index=False).agg(
    foo_count = ('foo', 'count'),
    bar_count = ('bar', 'count')
)

max_foo = grouped_df['foo_count'].max()
max_bar = grouped_df['bar_count'].max()

target_rows = grouped_df[(grouped_df['foo_count'] == max_foo) | (grouped_df['bar_count'] == max_bar)]

如果你的实际统计列是输出中显示的GDP、population，把上述代码里的foo、bar替换成对应的列名即可。

内容的提问来源于stack exchange，提问作者four-eyes