如何从Pandas DataFrame中按性别分层抽取指定数量样本？

原代码存在的问题 #

• 数据集拷贝错误：DD=X是浅拷贝，修改DD会直接改动原数据集X，应该用DD = X.copy()创建独立副本
• 循环逻辑失效：用range(len(DD))循环时，DD的长度会随drop操作动态变化，导致后续索引越界或部分数据未被处理
• 特征名混淆：代码中用DD.race[i]判断性别，但你的数据集里性别特征是gender，属于明显的笔误
• 抽样逻辑错误：
  • 第一个条件仅当男性数量等于目标值时才删除当前男性，无法精准保留maleSize条男性样本
  • 第二个条件判断女性数量超标时删除当前女性，但这种顺序删除的方式不是随机抽样，不符合需求
• 运行效率极低：循环遍历48000条数据并反复执行drop操作，Pandas的drop属于高开销操作，会导致代码运行极慢

正确实现方式 #

直接拆分男女子集，用Pandas内置的sample方法随机抽样后合并，既高效又能满足随机抽样的要求：

import pandas as pd

def Sample(X, maleSize, femaleSize):
    # 拆分男女样本子集
    male_df = X[X['gender'] == 1]
    female_df = X[X['gender'] == 0]
    
    # 提前检查样本数量是否足够
    if len(male_df) < maleSize:
        raise ValueError(f"男性样本仅{len(male_df)}条，不足要求的{maleSize}条")
    if len(female_df) < femaleSize:
        raise ValueError(f"女性样本仅{len(female_df)}条，不足要求的{femaleSize}条")
    
    # 分别随机抽取目标数量的样本
    sampled_male = male_df.sample(n=maleSize, random_state=42)  # random_state可选，用于固定抽样结果
    sampled_female = female_df.sample(n=femaleSize, random_state=42)
    
    # 合并样本并打乱顺序，避免结果集中男女分组排列
    Y = pd.concat([sampled_male, sampled_female]).sample(frac=1, random_state=42)
    
    return Y

代码说明 #

1. 用布尔索引精准拆分男女数据集，确保只从对应群体中抽样
2. sample方法是Pandas专门为随机抽样设计的函数，比手动循环高效数十倍
3. 合并后打乱顺序，保证结果集的随机性
4. 增加样本数量检查，避免因原数据样本不足导致的抽样失败

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