Pandas如何自定义DataFrame id分组下class列的0/1分布比例

核心逻辑 #

全量样本的class=0、class=1总量固定，id取值范围固定为0-50共51个分组，不要先全量随机分配id再调整——这种方式可调空间极小，应该先给需要自定义比值的分组预留配额，优先分配，剩余样本再随机分给剩下的分组，就能支持尽可能多的分组自定义比值。

具体操作步骤 #

• 先梳理自定义规则：列清楚需要调整比值的id、对应id的期望0/1比值（0样本数/1样本数）、单id分配的总样本量。
• 做配额可行性校验：逐个计算每个自定义id需要的0样本数、1样本数，公式为单id1样本数=round(单id总样本量/(比值+1))，单id0样本数=单id总样本量-单id1样本数。所有自定义id的0样本数总和不能超过全局总0样本量66932，1样本数总和不能超过全局总1样本量46337，超出的话就调小对应id的样本量，或者调整比值到合理范围。
• 优先分配自定义分组样本：把原数据集拆成class=0、class=1两个独立样本池，对每个自定义id，从两个池子里无放回抽取对应数量的样本，打上对应的id标签，抽中的样本直接从样本池里移除，避免重复分配。
• 分配剩余样本：两个样本池剩下的样本合并，随机分配给没有设置自定义规则的id，这部分分组的0/1比值会自然贴近全局1.444的水平，不会打乱自定义分组的比值。

最大化可调分组数量的技巧 #

• 自定义分组的单组样本量不要设置过大，单组样本量控制在1000-2000区间时，51个分组里最多可以支持45个以上的分组做自定义比值调整，只要总配额不超全局上限即可。
• 如果要设置极端比值（比如某分组全为0、全为1），尽量把这类分组的样本量设小（比如几百条），减少对全局配额的占用，留出更多空间给其他自定义分组。

参考实现代码 #

import pandas as pd
import numpy as np

# 初始化样本池，替换成你自己的数据集读取逻辑即可
# df = pd.read_csv('your_dataset.csv')
class0_pool = df[df['class'] == 0].reset_index(drop=True)
class1_pool = df[df['class'] == 1].reset_index(drop=True)

# 自定义规则格式：{id值: (目标0/1比值, 该id总样本量)}
custom_id_rules = {
    0: (3.2, 1600),
    1: (0.7, 1800),
    2: (2.5, 1200),
    # 按需补充更多自定义id规则
}
assigned_list = []
used_ids = set(custom_id_rules.keys())

# 先分配自定义id的样本
for cid, (target_ratio, sample_total) in custom_id_rules.items():
    n1 = round(sample_total / (target_ratio + 1))
    n0 = sample_total - n1
    # 无放回抽样
    s0 = class0_pool.sample(n=n0, replace=False)
    s1 = class1_pool.sample(n=n1, replace=False)
    # 打id标签
    s0['id'] = cid
    s1['id'] = cid
    assigned_list.append(pd.concat([s0, s1]))
    # 从样本池移除已分配样本
    class0_pool = class0_pool.drop(s0.index).reset_index(drop=True)
    class1_pool = class1_pool.drop(s1.index).reset_index(drop=True)

# 分配剩余样本
rest_samples = pd.concat([class0_pool, class1_pool]).reset_index(drop=True)
# 取出未做自定义设置的id列表
available_ids = [i for i in range(51) if i not in used_ids]
# 随机分配id
rest_samples['id'] = np.random.choice(available_ids, size=len(rest_samples))

# 合并得到最终结果
final_df = pd.concat(assigned_list + [rest_samples]).reset_index(drop=True)

运行时如果出现抽样数量超过池子剩余量的报错，说明自定义规则的总配额超了全局上限，对应调小自定义id的样本量或者调整比值即可。

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