Pandas：按列分组后选取前2个最大分组的实现方法

问题描述 #

原始DataFrame：

a   b
i1  t1
i1  t2
i2  t3
i2  t1
i2  t3
i3  t2

需要按列a分组后，提取行数最多的前2个分组的所有数据，最终结果要保留这两个分组的全部行，而不是每个分组的前2行。之前用df.groupby("a").head(2)无法实现，因为该方法只是取每个分组的前2行，不是筛选分组本身。

解决方案 #

方法1：分步实现 #

1. 先统计每个分组的行数：

group_counts = df.groupby('a').size()

2. 筛选出行数最多的前2个分组的名称：

top_2_groups = group_counts.nlargest(2).index

3. 过滤原DataFrame，只保留这些分组的所有行：

result_df = df[df['a'].isin(top_2_groups)]

方法2：链式简化写法 #

把上述步骤合并成一行代码，更简洁：

result_df = df[df['a'].isin(df.groupby('a').size().nlargest(2).index)]

可选：按分组行数排序 #

如果需要让结果按分组行数从多到少排列，可以在最后追加排序逻辑：

result_df = df[df['a'].isin(top_2_groups)].sort_values(
    'a', 
    key=lambda x: x.map(group_counts), 
    ascending=False
)

此时输出会和期望的顺序完全一致：

a   b
i2  t3
i2  t1
i2  t3
i1  t1
i1  t2

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