如何用Python实现分层随机分组 分配对照组与三个处理组且分布均衡

你的需求属于典型的分层随机抽样场景，核心思路是按Score字段的三个取值（high/medium/low）分成三个独立子集，每个子集内按照相同比例分配到四个分组，即可保证所有分组的评分分布完全一致。
分组比例要求：

• control组占总样本20%
• 三个treatment组各占总样本的 80% ÷ 3 ≈ 26.67%

依赖库：pandas、numpy，未安装可执行pip install pandas numpy安装。

import pandas as pd
import numpy as np

# 1. 加载数据，示例为模拟数据，实际使用替换为你的数据加载逻辑即可
# df = pd.read_csv("你的数据集路径.csv")
df = pd.DataFrame({
    "Score": ["high"]*300 + ["medium"]*300 + ["low"]*300 # 模拟900条样本，三类评分各占1/3
})

# 2. 定义分组占比规则
GROUP_RATIO = {
    "control": 0.2,
    "treatment 1": 0.8 / 3,
    "treatment 2": 0.8 / 3,
    "treatment 3": 0.8 / 3
}

# 3. 单层样本分组函数
def assign_group_for_layer(layer_df, ratio_rule):
    # 打乱当前层样本顺序
    shuffled_layer = layer_df.sample(frac=1, random_state=42).reset_index(drop=True)
    n_sample = len(shuffled_layer)
    # 计算每个组的样本量
    group_cnt = {group: int(np.round(ratio * n_sample)) for group, ratio in ratio_rule.items()}
    # 补全四舍五入导致的样本量差
    total_cnt = sum(group_cnt.values())
    if total_cnt < n_sample:
        group_cnt[list(ratio_rule.keys())[-1]] += (n_sample - total_cnt)
    # 生成分组标签
    group_labels = []
    for group, cnt in group_cnt.items():
        group_labels.extend([group] * cnt)
    shuffled_layer["group"] = group_labels
    return shuffled_layer

# 4. 按Score分层执行分组，合并结果
df_result = df.groupby("Score", group_keys=False).apply(assign_group_for_layer, ratio_rule=GROUP_RATIO)

# 可选：验证结果是否符合要求
if __name__ == "__main__":
    print("=== 整体分组占比 ===")
    print(df_result["group"].value_counts(normalize=True).round(4))
    print("\n=== 各分组内的评分分布 ===")
    for group in df_result["group"].unique():
        print(f"\n{group}：")
        print(df_result.query(f"group == '{group}'")["Score"].value_counts(normalize=True).round(4))

• 代码中的random_state=42是为了固定随机种子保证结果可复现，实际使用时可删除或修改为其他整数值。
• 兼容任意样本量的数据集，自动处理四舍五入导致的样本量偏差。
• 输出的df_result为增加了group分组标签的最终数据集，可直接导出使用。

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