Python实现将模型分数划分为分布一致的4个实验分组

要满足各组分数分布一致的要求，采用分层随机抽样即可实现：先将分数按区间分箱，每个箱内按指定比例分配组别，保证每个分数区间的样本在各组的占比与整体完全匹配。

依赖安装 #

你需要用到pandas和numpy两个常用数据处理库，如未安装可执行：
pip install pandas numpy

完整实现代码 #

import pandas as pd
import numpy as np

# ---------------------- 1. 读取/生成你的数据集 ----------------------
# 示例：生成10000条0-1分数的测试数据，替换为你自己的df读取逻辑即可
np.random.seed(42) # 固定随机种子保证结果可复现，不需要可以删掉
df = pd.DataFrame({'Score': np.random.uniform(0, 1, size=10000)})

# ---------------------- 2. 对分数分箱做分层依据 ----------------------
# 这里按分位数分20箱，数据量越大可以分越多箱，保证分布对齐精度越高
df['score_bin'] = pd.qcut(df['Score'], q=20, duplicates='drop')

# ---------------------- 3. 每个分箱内按比例分配组别 ----------------------
# 定义分组和对应权重：control占20%，三个treatment各占80%/3 ≈26.67%
group_labels = ['control', 'treatment 1', 'treatment 2', 'treatment 3']
group_weights = [0.2, 0.8/3, 0.8/3, 0.8/3]

def assign_group(bin_df):
    # 每个分箱内按权重随机分配组别
    return np.random.choice(group_labels, size=len(bin_df), p=group_weights)

df['group'] = df.groupby('score_bin', group_keys=False).apply(assign_group)

# ---------------------- 可选：验证分组结果 ----------------------
# 查看各组样本量占比
print("各组样本量占比：")
print(df['group'].value_counts(normalize=True).round(4)*100)

# 查看各组分数分布统计，均值、中位数、分位数应该基本一致
print("\n各组分数统计：")
print(df.groupby('group')['Score'].agg(['mean', 'median', 'min', 'max']).round(4))

说明 #

• 如果你的数据集本身已经有其他需要对齐的特征（比如用户年龄、地域），可以把这些特征也加入分箱维度，一起做分层抽样即可
• 分箱数量可以根据你的数据量调整：万级数据分10-20箱即可，十万级以上可以分50-100箱，分布对齐效果更好
• 样本量较大的情况下，分组占比误差会小于0.1%，分布一致性完全满足实验要求

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