如何用Pandas优雅选取具有唯一A-B组合的DataFrame行？

需求说明 #

现有包含A、B、C、D列的DataFrame，A列粒度最粗，B列粒度更细。需要为每个唯一的A、B值组合选取一行（示例中取每组的第一行）。

示例输入数据 #

import pandas as pd

df = pd.DataFrame({
    'A': ['a1', 'a1', 'a1', 'a1', 'a2', 'a2', 'a2', 'a2'],
    'B': ['b1', 'b1', 'b2', 'b2', 'b3', 'b3', 'b4', 'b4'],
    'C': ['c1', 'c2', 'c3', 'c4', 'c5', 'c6', 'c7', 'c8'],
    'D': ['d1', 'd2', 'd3', 'd4', 'd5', 'd6', 'd7', 'd8']
})

原始DataFrame输出：

A    B    C    D
0   a1   b1   c1   d1
1   a1   b1   c2   d2
2   a1   b2   c3   d3
3   a1   b2   c4   d4
4   a2   b3   c5   d5
5   a2   b3   c6   d6
6   a2   b4   c7   d7
7   a2   b4   c8   d8

期望结果 #

A    B    C    D
0   a1   b1   c1   d1
2   a1   b2   c3   d3
4   a2   b3   c5   d5
6   a2   b4   c7   d7

解决方案 #

以下是几种简洁高效的实现方式：

方法1：groupby + first() #

按A、B列分组后直接取每组第一行，是最直观的实现方式：

result = df.groupby(['A', 'B'], as_index=False).first()

• as_index=False 保持原列结构，不会将分组列设为索引
• first() 保留每组最先出现的行的所有列值

方法2：drop_duplicates #

利用去重逻辑，指定按A、B列判断重复并保留第一行：

result = df.drop_duplicates(subset=['A', 'B'], keep='first')

• subset=['A', 'B'] 指定去重判断的列
• keep='first' 保留每组第一行，若需最后一行可改为keep='last'

方法3：groupby + nth(0) #

适合需要灵活选取组内指定位置行的场景，这里取第0行（即第一行）：

result = df.groupby(['A', 'B'], as_index=False).nth(0)

• nth(0) 精确选取组内第0个位置的行，支持指定多个位置（如nth([0,2])取第1、3行）

方法4：groupby + head(1) #

返回每组的前N行，这里取1行：

result = df.groupby(['A', 'B']).head(1).reset_index(drop=True)

• reset_index(drop=True) 重置结果索引，避免原索引混乱

以上方法均可满足需求：drop_duplicates在大数据集上性能更优，groupby系列方法则更灵活，适合后续扩展分组操作的场景。

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