基于多条件为pandas主表添加Myset字段实现训练测试集拆分

以下代码完全满足3项分配规则，可直接运行：

• 已匹配参考表的MasterID直接沿用参考表的Myset取值
• 未匹配的MasterID以MyId为排序依据整组分配，保证同ID归属唯一集合
• 分配时动态平衡两组样本量，尽可能实现拆分后样本数均等

import pandas as pd

# 1. 左关联参考表，先完成已知ID的Myset赋值
df = df.merge(ref, on="MasterID", how="left")

# 2. 提取所有未匹配参考表的唯一MasterID，按对应MyId排序保证拆分可复现
unmatched_id = (
    df[df["Myset"].isna()]
    .groupby("MasterID", as_index=False)["MyId"]
    .min()
    .sort_values("MyId", ascending=True)
    .reset_index(drop=True)
)

# 3. 统计当前已分配的两组样本量，动态给未匹配ID整组分配集合
train_cnt = df[df["Myset"] == "TRAIN"].shape[0]
test_cnt = df[df["Myset"] == "TEST"].shape[0]
assign_res = []

for _, row in unmatched_id.iterrows():
    # 统计当前MasterID在主表中的总行数，整组分配到样本量更小的集合
    id_row_num = df[df["MasterID"] == row["MasterID"]].shape[0]
    if train_cnt <= test_cnt:
        assign_res.append("TRAIN")
        train_cnt += id_row_num
    else:
        assign_res.append("TEST")
        test_cnt += id_row_num

# 生成未匹配ID的分组映射表
unmatched_id["Myset"] = assign_res
full_map = pd.concat([ref, unmatched_id[["MasterID", "Myset"]]], ignore_index=True)

# 4. 回填所有Myset值，避免同ID跨集合问题
df = df.drop(columns=["Myset"]).merge(full_map, on="MasterID", how="left")

# 5. 按要求拆分数据集
train = df[df["Myset"] == "TRAIN"]
test = df[df["Myset"] == "TEST"]

校验提示：执行完可通过df.groupby("MasterID")["Myset"].nunique().max()校验，返回值为1即证明不存在同MasterID归属两个集合的问题；通过train.shape[0]和test.shape[0]可查看两组样本量差值，该方案会尽可能将差值控制在单个MasterID对应的最大行数范围内。

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