如何按Salesman分组，根据对应sales_limit筛选指定行数的DataFrame

需求说明 #

给定包含销售、客户信息的DataFrame，需为每个salesman保留最多n行数据，其中n是该销售对应的sales_limit值（该限额字典每周动态更新）。示例规则：

• salesman A、B的限额为2，各保留前2行
• salesman C的限额为4，大于现有行数，全部保留

示例数据 #

import pandas as pd

# 原始业务数据
df = pd.DataFrame({
    'salesman': ['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C'],
    'client_id': ['abc', 'def', 'ghi', 'jkl', 'mno', 'pqr', 'stu', 'vwx', 'yz'],
    'sales_amount': [100, 200, 150, 300, 250, 180, 400, 350, 500]
})

# 每周动态更新的销售限额字典
sales_limit = {'A': 2, 'B': 2, 'C': 4}

实现代码 #

方法1：GroupBy + Apply（逻辑直观） #

适合小数据量或需要灵活扩展筛选逻辑的场景：

def filter_group(group, limit_map):
    # 获取当前分组销售的限额，默认0（无数据时返回空）
    limit = limit_map.get(group.name, 0)
    return group.head(limit)

# 分组筛选后重置索引，还原成标准DataFrame结构
filtered_df = df.groupby('salesman').apply(filter_group, limit_map=sales_limit).reset_index(drop=True)

方法2：Cumcount + 映射（高效适配大数据） #

避免apply的性能损耗，适合百万级以上大规模数据：

# 为每个销售的行按顺序编号（从1开始计数）
df['row_seq'] = df.groupby('salesman').cumcount() + 1
# 映射每个销售对应的限额
df['limit'] = df['salesman'].map(sales_limit)
# 筛选行号≤限额的记录，清理临时辅助列
filtered_df = df[df['row_seq'] <= df['limit']].drop(columns=['row_seq', 'limit']).reset_index(drop=True)

结果验证 #

运行代码后，filtered_df将精准符合需求：移除A的第3行（client_id=ghi）和B的第3行（client_id=pqr），保留C的全部行。

内容的提问来源于stack exchange，提问作者Bruce Wayne 0005