如何基于列值从DataFrame中随机采样生成平衡数据集

需求背景 #

现有名为df的DataFrame，其Priority列的样本分布如下：

df.Priority.value_counts()

P3    39506
P2    3038 
P4    1138 
P1    1117 
P5    252  
Name: Priority, dtype: int64

需要生成平衡数据集df_balanced，将P3类别的样本数限制为5000，其余类别保留全部样本，最终分布如下：

P3    5000
P2    3038 
P4    1138 
P1    1117 
P5    252  
Name: Priority, dtype: int64

问题修正与解决方案 #

你之前的代码中，s0.union(df)是错误的——df是DataFrame对象，并非索引集合，无法直接和P3的采样索引合并。以下是两种可行的实现方式：

方法1：分块合并（直观易懂）

import pandas as pd

# 对P3类别随机采样5000行（random_state可选，用于固定采样结果）
p3_sampled = df[df['Priority'] == 'P3'].sample(n=5000, random_state=42)
# 提取所有非P3类别的行
non_p3_data = df[df['Priority'] != 'P3']
# 合并两部分数据并重置索引
df_balanced = pd.concat([p3_sampled, non_p3_data]).reset_index(drop=True)

方法2：索引合并（更高效）

# 获取P3采样行的索引
p3_idx = df[df['Priority'] == 'P3'].sample(5000).index
# 获取所有非P3行的索引
non_p3_idx = df[df['Priority'] != 'P3'].index
# 通过合并索引提取目标数据并重置索引
df_balanced = df.loc[p3_idx.union(non_p3_idx)].reset_index(drop=True)

结果验证 #

运行以下代码即可确认样本分布符合预期：

print(df_balanced['Priority'].value_counts())

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