如何在Pandas中按col1分层随机抽取指定数量样本？

原始DataFrame创建代码 #

import pandas as pd
import numpy as np

ds = {'col1' : [1,1,1,1,1,1,1,2,2,2,2,3,3,3,3,3,4,4,4,4,4,4,4,4,4],
      'col2' : [12,3,4,5,4,3,2,3,4,6,7,8,3,3,65,4,3,2,32,1,2,3,4,5,32],
      }

df = pd.DataFrame(data=ds)

抽样需求 #

• 抽取3条col1 == 1的随机记录
• 抽取2条col1 == 2的随机记录
• 抽取2条col1 == 3的随机记录
• 抽取3条col1 == 4的随机记录

解决方案 #

可以利用Pandas的groupby().sample()方法，通过字典指定每个分组的抽样数量，一步完成分组抽样：

# 定义每个分组的抽样数量
sample_counts = {1:3, 2:2, 3:2, 4:3}

# 按col1分组，按指定数量抽样
sampled_df = df.groupby('col1', group_keys=False).sample(n=sample_counts)

# 查看结果
print(sampled_df)

代码说明 #

• groupby('col1')：按照col1列的值对DataFrame进行分组
• group_keys=False：避免结果中保留分组键作为索引的一部分，保持原始DataFrame的索引结构
• sample(n=sample_counts)：传入字典sample_counts，指定每个分组对应的抽样数量，Pandas会自动为每个组抽取对应数量的随机样本

如果需要固定抽样结果（保证可复现），可以添加random_state参数：

sampled_df = df.groupby('col1', group_keys=False).sample(n=sample_counts, random_state=42)

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