如何基于Grade和Category两列对DataFrame进行分层抽样？

创建初始DataFrame #

# Create a dictionary of students
students = {
    'Name': ['Lisa', 'Kate', 'Ben', 'Kim', 'Josh',
             'Alex', 'Evan', 'Greg', 'Sam', 'Ella'],
    'ID': ['001', '002', '003', '004', '005', '006',
           '007', '008', '009', '010'],
    'Grade': ['A', 'A', 'C', 'B', 'B', 'B', 'C',
              'A', 'A', 'A'],

    'Category': [1, 1, 1, 1, 0, 0, 0, 0, 1, 0]
}

# Create dataframe from students dictionary
df = pd.DataFrame(students)

# view the dataframe
df

原数据分布 #

• Grade：A占50%，B占30%，C占20%
• Category：1占50%，0占50%

需求 #

对上述DataFrame进行分层抽样，确保抽样后的样本同时保持Grade和Category的原分布比例。

已尝试的方法及问题 #

1. 仅按Grade分组抽样

df.groupby('Grade', group_keys=False).apply(lambda row: row.sample(frac=0.6))

仅能保证Grade的分布符合要求，但无法维持Category的原比例。

2. 按Grade和Category联合分组抽样

df.groupby(['Grade','Category'],group_keys=False).apply(lambda row: row.sample(frac=0.6))

该方法会对每个交叉分组（如A&1、A&0等）单独按比例抽样，最终整体的Category或Grade比例可能偏离原分布，不符合预期。

正确实现方案 #

要同时保证两列的分布比例，需要采用双重分层控制：先按其中一列（比如Grade）分层，在每个分层内再按另一列（Category）的全局原比例分配抽样数量。

代码实现

# 先获取全局Category的分布比例
category_dist = df['Category'].value_counts(normalize=True)

# 按Grade分组，在每个组内按Category的全局比例抽样
sampled_df = df.groupby('Grade', group_keys=False).apply(
    lambda grade_group: grade_group.groupby('Category', group_keys=False).apply(
        lambda cat_group: cat_group.sample(frac=0.6 * category_dist[cat_group.name])
    )
).reset_index(drop=True)

另一种直观实现（按计数分配）

如果需要更精准的样本数量控制，可以先计算每个Grade分组的总抽样数，再在组内按Category比例分配：

# 计算每个Grade分组的总抽样数（按总体60%比例）
grade_sample_counts = (df['Grade'].value_counts(normalize=True) * len(df) * 0.6).round().astype(int)

# 按Grade分组，在组内按Category比例分配抽样数
sampled_df = pd.concat([
    df[df['Grade'] == grade].groupby('Category', group_keys=False).apply(
        lambda cat_group: cat_group.sample(n=int(grade_sample_counts[grade] * category_dist[cat_group.name]))
    )
    for grade in grade_sample_counts.index
]).reset_index(drop=True)

验证结果 #

可以通过以下代码确认抽样后的分布是否符合要求：

# 验证Grade分布
print("抽样后Grade分布：")
print(sampled_df['Grade'].value_counts(normalize=True))

# 验证Category分布
print("\n抽样后Category分布：")
print(sampled_df['Category'].value_counts(normalize=True))

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