如何在Pandas DataFrame分组内按键序选取前N个SEQ的前2行

问题描述 #

给定如下数据集：

import pandas as pd

data = [
    ('A', 'X'),
    ('A', 'X'),
    ('A', 'Y'),
    ('A', 'Z'),
    ('B', 1),
    ('B', 1),
    ('B', 2),
    ('B', 2),
    ('B', 3),
    ('B', 3),
    ('C', 'L-7'),
    ('C', 'L-9'),
    ('C', 'L-9'),
    ('T', 2020),
    ('T', 2020),
    ('T', 2025)
]

df = pd.DataFrame(data, columns=['ID', 'SEQ'])

需求：

1. 按ID分组，在每个ID组内按原数据顺序选取前2个不同的SEQ值
2. 对每个选中的SEQ，保留其对应的最多前2行（若存在）

例如ID为A的分组中，前2个不同的SEQ是X和Y，因此保留X的2行、Y的1行，排除Z的行。

尝试使用df.groupby(['ID','SEQ']).head(2)未达到效果，该代码仅保留每个ID-SEQ组的前2行，但未过滤掉每个ID内超出前2个的SEQ。

解决方案（链式调用风格） #

使用Pandas链式调用实现需求，代码如下：

result = (df
          # 为每个ID内的SEQ标记首次出现的顺序（从0开始）
          .assign(seq_order=lambda x: x.groupby('ID')['SEQ'].transform(lambda s: s.factorize()[0]))
          # 筛选每个ID内前2个出现的SEQ
          .query('seq_order < 2')
          # 对每个ID-SEQ组保留最多前2行
          .groupby(['ID', 'SEQ'])
          .head(2)
          # 移除临时辅助列
          .drop('seq_order', axis=1)
          # 重置索引（可选，按需调整）
          .reset_index(drop=True)
)

print(result)

代码解释 #

1. assign(seq_order=...)：通过factorize()为每个ID分组内的SEQ标记首次出现的顺序，返回的元组第一个元素是顺序编号（从0开始），这样每个ID内前2个不同的SEQ对应编号0和1。
2. query('seq_order < 2')：筛选出每个ID内前2个出现的SEQ的所有行，直接排除后续出现的SEQ（比如ID A中的Z）。
3. groupby(['ID', 'SEQ']).head(2)：确保每个ID-SEQ组合最多保留前2行，避免单个SEQ的行数超过2的情况。
4. drop('seq_order', axis=1)：移除临时生成的辅助列，得到干净的结果。

最终输出 #

运行代码后得到的结果如下：

ID    SEQ
0   A      X
1   A      X
2   A      Y
3   B      1
4   B      1
5   B      2
6   B      2
7   C    L-7
8   C    L-9
9   C    L-9
10  T   2020
11  T   2020
12  T   2025

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