如何按条件筛选Pandas DataFrame行：保留各类别Top3成员

需求说明 #

现有一个已按类别分组、并依据分数排序的Pandas DataFrame，需要保留每个类别中最多前3个成员；若某类别成员数量不足3个，则保留该类别全部成员。

原数据与初始DataFrame #

import pandas as pd

data = [['Tom', 'A', 10], ['Jack', 'A', 9], ['Andy', 'A', 8],
        ['Lily', 'A', 7], ['Johan', 'B', 10], ['Ewa', 'B', 9]]

df = pd.DataFrame(data, columns=['Name', 'Category', 'Score'])

原DataFrame表格：

解决方案 #

直接使用groupby()按Category分组，再调用head(3)即可实现需求：

# 保留每个类别最多前3个成员
result_df = df.groupby('Category').head(3)

结果验证 #

执行上述代码后得到的DataFrame如下：

原理说明 #

• groupby('Category')将DataFrame按类别分组；
• head(3)会为每个分组提取前3行数据，若分组内的行数不足3，则直接保留该分组的全部行，完全匹配需求。

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