基于Pandas按n_c列值筛选每组指定数量记录的需求

需求说明 #

创建新的DataFrame，按id和sys_id分组（每组内n_c值统一），保留每组中对应n_c列值数量的前若干条记录：

• n_c=1时保留每组前1条
• n_c=2时保留每组前2条
• n_c=4时保留每组前4条
  其余记录直接剔除。

输入数据 #

预期输出 #

初始代码 #

import pandas as pd
  
# 初始化列表数据
data = {'id': ['G400', 'G406', 'G406', 'G406','G408','G408', 'G408', 'G408','G408'],
        'sys_id': ['02f4', '993c', '993c', '993c', '996d','996d', '996d', '996d','996d'],
        'n_c': [1,2,2,2,4,4,4,4,4],
        'extract_id':[284,389,382,129,396,378,368,333,222]}  
  
# 创建DataFrame
df = pd.DataFrame(data)

解决方案 #

通过groupby分组后，结合apply和head方法实现自定义筛选：

# 按id和sys_id分组，每组保留前n_c条记录（n_c为组内统一值）
filtered_df = df.groupby(['id', 'sys_id'], group_keys=False).apply(
    lambda x: x.head(x['n_c'].iloc[0])
)

# 查看结果
print(filtered_df)

代码说明 #

1. 分组逻辑：以id和sys_id作为分组依据，保证每组的n_c值唯一匹配需求
2. 筛选逻辑：在apply的匿名函数中，取每组第一个n_c值作为head的参数，保留对应数量的前N条记录
3. 结果格式：设置group_keys=False避免分组键被添加为索引，保持结果与原DataFrame格式一致

运行后得到的filtered_df完全符合预期输出。

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