如何在Python Pandas中按指定类别比例随机抽取数据样本？

核心思路 #

先将原数据按TG列拆分为两个子集，分别抽取对应数量的样本后合并，最终得到符合比例要求的数据集。

具体代码 #

import pandas as pd

# 假设你的原DataFrame名为df
# 1. 拆分TG=0和TG=1的子集
df_tg0 = df[df['TG'] == 0]
df_tg1 = df[df['TG'] == 1]

# 2. 设置总样本量和TG=1的目标占比
total_samples = 200000
tg1_target_ratio = 0.05

# 3. 计算各子集需要抽取的数量
tg1_sample_num = int(total_samples * tg1_target_ratio)
tg0_sample_num = total_samples - tg1_sample_num

# 4. 随机抽样（replace=False表示不重复抽样，原数据量足够的情况下使用）
sampled_tg1 = df_tg1.sample(n=tg1_sample_num, random_state=42)  # random_state固定可复现结果
sampled_tg0 = df_tg0.sample(n=tg0_sample_num, random_state=42)

# 5. 合并两个抽样后的子集，并打乱顺序
final_sample = pd.concat([sampled_tg1, sampled_tg0]).sample(frac=1, random_state=42).reset_index(drop=True)

# 验证比例
print(f"TG=1占比: {final_sample['TG'].mean():.2%}")

注意事项 #

• 如果原数据中TG=1的总数量小于tg1_sample_num，可以将sample方法的replace参数设为True（即有放回抽样），但你的场景中原TG=1有3万条，远大于需要的1万条，无需此操作。
• 设置random_state是为了让抽样结果可复现，若不需要固定结果可以去掉该参数。

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