如何用Pandas从多类别列数据中抽取近似等量行数的样本？

针对从含多分类列的大数据集（1亿行）中抽取指定样本量（20万行），且要求各分类列的每个类别样本数尽可能均衡（±5%误差可接受）的需求，以下是两种高效可行的实现方案，支持replace=True参数：

方案1：加权抽样（高效适配大数据场景） #

原理 #

通过给每个样本分配与所在分类稀缺性挂钩的权重，让占比低的分类样本拥有更高的抽样概率，最终实现各分类样本量的均衡。该方案无需分组循环，底层由Pandas优化，处理1亿行数据效率极高。

代码实现 #

import pandas as pd
import numpy as np

# 计算各分类列的类别占比（归一化）
cat1_ratio = df1['cat1'].value_counts(normalize=True)
cat2_ratio = df1['cat2'].value_counts(normalize=True)
cat3_ratio = df1['cat3'].value_counts(normalize=True)

# 为每个样本计算权重：各分类占比的倒数乘积，稀缺分类的样本权重更高
df1['sample_weight'] = (1 / df1['cat1'].map(cat1_ratio)) * \
                       (1 / df1['cat2'].map(cat2_ratio)) * \
                       (1 / df1['cat3'].map(cat3_ratio))

# 基于权重进行随机抽样，支持replace=True
sample = df1.sample(n=200000, replace=True, weights='sample_weight', random_state=123)

# 可选：删除临时权重列
sample.drop('sample_weight', axis=1, inplace=True)

方案2：交叉分组抽样（均衡性更精准） #

原理 #

按所有分类列的交叉组合进行分组，给每个组合分配近似相等的抽样数量，最终保证单一分类维度的样本量均衡。适合对均衡性要求更高的场景，若交叉组合数量不多（如你的场景中2×3×4=24组），大数据下也可高效运行。

代码实现 #

total_sample = 200000

# 按三个分类列的交叉组合分组
grouped = df1.groupby(['cat1', 'cat2', 'cat3'])
num_groups = len(grouped)

# 计算每个分组的基础抽样数，剩余样本数分配给前N个分组
per_group = total_sample // num_groups
remaining = total_sample % num_groups

sample_list = []
for idx, (group_key, group_data) in enumerate(grouped):
    # 确定当前分组的抽样数量
    n = per_group + 1 if idx < remaining else per_group
    # 若分组内样本不足，启用replace=True补充
    sample_list.append(group_data.sample(n=n, replace=len(group_data) < n, random_state=123))

# 合并所有分组的抽样结果
sample = pd.concat(sample_list, ignore_index=True)

均衡性验证（修正原assert逻辑） #

原验证代码存在语法错误，以下是符合需求的验证方式，确保各分类样本量误差在±5%以内：

total_sample = 200000

# 验证cat1（2个类别）
cat1_counts = sample['cat1'].value_counts()
assert abs(cat1_counts['a'] - cat1_counts['b']) / total_sample <= 0.05, "cat1均衡性不达标"

# 验证cat2（3个类别）
target_cat2 = total_sample / 3
for cnt in sample['cat2'].value_counts().values:
    assert abs(cnt - target_cat2) / target_cat2 <= 0.05, "cat2均衡性不达标"

# 验证cat3（4个类别）
target_cat3 = total_sample / 4
for cnt in sample['cat3'].value_counts().values:
    assert abs(cnt - target_cat3) / target_cat3 <= 0.05, "cat3均衡性不达标"

大数据场景优化提示 #

• 优先选择方案1，加权抽样的内存占用和计算效率远高于分组抽样，适合1亿行级别的数据集。
• 若使用方案2，避免过多的交叉分类组合（如超过1000组），否则分组操作会显著变慢。
• replace=True仅在分组内样本量不足目标抽样数时自动生效，无需额外判断。

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