如何对DataFrame进行多级分组并筛选各子组内的Top N记录？

先给出示例数据的初始化代码：

import pandas as pd

df = pd.DataFrame({
    'Date': ['07-08-2023','06-08-2023','08-08-2023','08-08-2023','07-08-2023','06-08-2023','06-08-2023','08-08-2023'],
    'Category': ['A','A','A','A','B','B','B','B'],
    'Products': ['abc','xyz','abc','def','nop','hjk','qrs','qrs'],
    'Sales': [1,2,4,3,4,5,2,3]
})

下面是两种不用循环的优雅实现方式：

方法一：用辅助列排名筛选 #

通过transform计算每组产品的总销售额，再给分类内的产品排名，最后筛选Top N：

N = 2

# 给每条记录添加对应产品在分类内的总销售额
df['total_sales'] = df.groupby(['Category', 'Products'])['Sales'].transform('sum')
# 按分类对总销售额降序排名，dense方法避免并列排名导致漏选
df['rank'] = df.groupby('Category')['total_sales'].rank(method='dense', ascending=False)

# 筛选排名前N的产品记录，去掉辅助列
result = df[df['rank'] <= N].drop(columns=['total_sales', 'rank'])
print(result)

方法二：聚合筛选后合并 #

先找出每个分类的Top N产品列表，再和原表关联获取完整记录：

N = 2

# 先聚合计算每个产品的总销售额，再按分类取Top N产品
top_products = df.groupby(['Category', 'Products'])['Sales'].sum() \
                .groupby('Category', group_keys=False) \
                .nlargest(N) \
                .reset_index()[['Category', 'Products']]

# 合并回原表，得到仅包含Top N产品的所有记录
result = pd.merge(df, top_products, on=['Category', 'Products'])
print(result)

两种方法都能高效完成需求，不用写循环，适合处理大规模数据。

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