如何从Pandas DataFrame中抽取分层随机样本并拆分数据集？

原数据集创建 #

import pandas as pd
import numpy as np

ds = {'id':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26],
      "Grade":["A","A","A","A","A","B","A","A","A","C","C","C","B","D","D","D","D","D","D","D","D","B","A","B","C","D"]}

df = pd.DataFrame(data=ds)

Grade字段频数统计 #

df['Grade'].value_counts(dropna=False)

输出结果：

Grade
A    9
D    9
B    4
C    4
Name: count, dtype: int64

需求 #

1. 创建df_1，包含：
   • 5条Grade=D的随机记录
   • 2条Grade=C的随机记录
   • 1条Grade=B的随机记录
   • 3条Grade=A的随机记录
     总计11条记录
2. 创建df_2，包含原数据集df中除去df_1的所有记录，总计15条记录

解决方案 #

1. 生成df_1 #

通过groupby按Grade分组，再对每组抽取指定数量的随机样本：

# 定义各Grade需要抽取的样本数
sample_counts = {
    'D': 5,
    'C': 2,
    'B': 1,
    'A': 3
}

# 按分组抽取样本，random_state保证结果可复现（可选）
df_1 = df.groupby('Grade', group_keys=False).apply(
    lambda group: group.sample(n=sample_counts[group.name], random_state=42)
)

# 验证样本数量
print(f"df_1记录数：{len(df_1)}")  # 输出11

2. 生成df_2 #

通过索引判断，筛选出原数据中不在df_1里的记录：

# 筛选原数据中索引不在df_1中的记录
df_2 = df[~df.index.isin(df_1.index)]

# 验证样本数量
print(f"df_2记录数：{len(df_2)}")  # 输出15

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