如何对pandas DataFrame按分组列的组占比加权抽取指定总量样本

实现方案 #

你可以先按分组原始数据量计算加权抽样配额，再按配额分组抽样，保证总样本量为18000的同时符合占比要求，完整实现代码如下：

import pandas as pd

# 定义总抽样量和随机种子
TOTAL_SAMPLE = 18000
RANDOM_STATE = 1

# 统计各prefix分组的原始数据量
group_sizes = df.groupby("prefix").size()
# 按分组占比计算对应抽样数，四舍五入取整
sample_n_per_group = (group_sizes / group_sizes.sum() * TOTAL_SAMPLE).round().astype(int)

# 修正四舍五入带来的总样本量偏差，调整量作用在最大分组上避免影响小分组抽样
count_diff = TOTAL_SAMPLE - sample_n_per_group.sum()
if count_diff != 0:
    sample_n_per_group[sample_n_per_group.idxmax()] += count_diff

# 按计算出的配额分组抽样
sample_df = df.groupby("prefix", group_keys=False).apply(
    lambda group: group.sample(n=sample_n_per_group.loc[group.name], random_state=RANDOM_STATE)
).reset_index(drop=True)

注意事项 #

• 若存在分组计算出的抽样数为0的情况，可以在计算sample_n_per_group后添加判断逻辑，给小分组设置最低抽样数（比如至少抽1条），再重新修正总样本量即可
• 若使用pandas 1.1.0及以上版本，也可直接调用groupby.sample方法的weights参数实现加权，但上述手动计算配额的方式对总样本量的控制更精准，也更易调试

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