如何在Pandas中按两列分组后返回数据集的Top N值？

已完成的前置操作 #

• 创建示例DataFrame：

# 创建随机DataFrame
dctf =  {"A":["First","First","First","First","First","First","Second","Second","Second","Second","Second","Third","Third","Third","Third","Third"],
      "B":["one","two","one","two","one","two","one","two","three","one","two","three","three","one","two","three"],
      "C":["Apple","Orange","Leaf","foo","Leaf","Apple","foo","Leaf","Apple","Apple","Orange","Orange","Orange","foo","foo","Leaf"]
      }
df_ex = pd.DataFrame(dctf)

• 按["A","B"]分组计算C列的唯一值数量，生成counts列：

# 生成counts列
df_ex['counts']=df_ex.groupby(["A","B"])["C"].transform('nunique')

• 再次按["A","B"]分组，获取每组的counts最大值：

# 分组获取每组counts最大值
df_ex = df_ex.groupby(["A","B"], as_index=False)['counts'].max()

执行后得到的中间结果：

A       B       counts
0   First   one     2
1   First   two     3
2   Second  one     2
3   Second  three   1
4   Second  two     2
5   Third   one     1
6   Third   three   2
7   Third   two     1

需求 #

按A列分组后，返回每组中counts的最大值（或Top 5值），期望结果如下：

A       B       counts
1   First   two     3
2   Second  one     2
4   Second  two     2
6   Third   three   2

解决方案 #

1. 获取每组counts最大值对应的行 #

方法一：分组取最大值后匹配筛选

# 按A分组获取每组的最大counts值
max_counts = df_ex.groupby('A')['counts'].max().reset_index(name='max_count')
# 匹配筛选出符合条件的行
result_max = df_ex.merge(max_counts, on='A').query('counts == max_count').drop(columns='max_count')
print(result_max)

方法二：使用rank排名筛选

# 按A分组对counts降序排名（dense方法确保相同值排名一致）
df_ex['rank'] = df_ex.groupby('A')['counts'].rank(ascending=False, method='dense')
# 筛选排名为1的行（即最大值对应的行）
result_max = df_ex[df_ex['rank'] == 1].drop(columns='rank')
print(result_max)

2. 获取每组counts的Top N值（以Top5为例） #

基于rank方法调整筛选条件即可：

N = 5
# 按A分组对counts降序排名
df_ex['rank'] = df_ex.groupby('A')['counts'].rank(ascending=False, method='dense')
# 筛选排名<=N的行
result_topN = df_ex[df_ex['rank'] <= N].drop(columns='rank')
print(result_topN)

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