pandas按两列groupby后随机选组且col1每值对应唯一组的实现

需求说明 #

对df.groupby(['col1', 'col2'])生成的所有子分组进行随机选择，要求col1的每个取值仅对应1个被选中的col2子分组，返回所有被选中子分组的全部行数据。

示例输入数据 #

import pandas as pd    
df = pd.DataFrame({'col1': list('aaabbbabababbaaa'), 'col2': list('cdddccdsssssddcd'), 'val': range(0, 16)})

原df输出如下：

col1 col2  val
0     a    c    0
1     a    d    1
2     a    d    2
3     b    d    3
4     b    c    4
5     b    c    5
6     a    d    6
7     b    s    7
8     a    s    8
9     b    s    9
10    a    s   10
11    b    s   11
12    b    d   12
13    a    d   13
14    a    c   14
15    a    d   15

Pythonic实现方案 #

最简实现（merge匹配法，性能最优可读性最高） #

# 按col1分组后随机抽取1个col2值，去重后和原表关联匹配
selected_keys = df[['col1', 'col2']].groupby('col1').sample(n=1).drop_duplicates()
result = df.merge(selected_keys, on=['col1', 'col2'])

链式调用实现（单句完成需求） #

result = df.groupby('col1', group_keys=False).apply(
    lambda g: g[g['col2'] == g['col2'].sample(n=1).iloc[0]]
)

输出示例 #

每次运行结果随机，符合抽样要求：

col1 col2  val
4     b    c    4
5     b    c    5
8     a    s    8
10    a    s   10

方案说明 #

两种实现都完全基于pandas原生API，无需手动维护索引、无需额外调用numpy打乱数组，也没有显式for循环，代码简洁易读：

• 核心逻辑先对每个col1分组，随机选取1个对应的col2值
• 再过滤原表保留所有匹配(col1, col2)组合的行
• 大样本量下优先选merge方案，性能明显优于逐行判断匹配的写法

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