如何按唯一ID分层拆分多标签熔解格式数据集

数据集情况 #

• 规模约(200000, 700)的多标签（4个非互斥标签）熔解格式DataFrame，示例结构如下：

• 每个客户对应多个月度快照，每个快照下最多4个标签的0/1状态；无法将数据透视还原，否则会出现缺失值导致模型训练失败。

核心需求 #

1. 拆分后每个客户的所有记录仅属于训练/验证/测试集其中一个，无交叉；
2. 三个集合的各标签0/1分布比例需尽可能一致（分层拆分），比如训练集各标签分布如下时，验证、测试集需同比例近似：

已实现的无分层拆分逻辑 #

通过sklearn.train_test_split按客户ID拆分，代码如下：

from sklearn.model_selection import train_test_split

train_size = 0.7
val_size = 0.15
test_size = 0.15

# 获取唯一客户ID列表
unique_cust_list = df["customer_id"].unique().tolist()
 
# 拆分ID列表为训练、验证、测试集
train_id_list, remainder = train_test_split(unique_cust_list, test_size= val_size + test_size)
val_id_list, test_id_list = train_test_split(remainder, test_size=test_size / (val_size + test_size))

# 过滤得到对应数据集
train_df = df[df["customer_id"].isin(train_id_list)]
val_df = df[df["customer_id"].isin(val_id_list)]
test_df = df[df["customer_id"].isin(test_id_list)]

待解决问题 #

sklearn的StratifiedShuffleSplit、StratifiedGroupKFold不原生支持这种多标签熔解格式的按客户ID分层拆分，需自定义分层逻辑实现标签分布均衡。

解决方案 #

步骤1：构建客户级分层键 #

为每个客户生成代表其标签分布的特征，作为分层依据，两种可选方式：

方式A：客户各标签的目标均值（适合多时间快照客户）

计算每个客户在各标签下的target均值（即1的占比），再离散化后拼接成分层键：

# 按客户+标签分组，计算target均值
cust_label_stats = df.groupby(["customer_id", "label"])["target"].mean().unstack(fill_value=0)

# 对均值离散化，避免连续值导致分层失效
def discretize(val):
    if val >= 0.7:
        return "high"
    elif val >= 0.3:
        return "mid"
    else:
        return "low"

# 生成客户分层键
cust_strat_key = cust_label_stats.applymap(discretize).agg("-".join, axis=1)

方式B：客户各标签的末次状态（适合时间快照少的客户）

取每个客户各标签的末次target值，拼接成分层键：

# 按客户+标签取末次target值
cust_label_last = df.groupby(["customer_id", "label"])["target"].last().unstack(fill_value=0)

# 生成客户分层键
cust_strat_key = cust_label_last.astype(str).agg("-".join, axis=1)

步骤2：基于分层键拆分客户ID #

用train_test_split以分层键为依据拆分，保证各集合标签分布一致：

from sklearn.model_selection import train_test_split

train_size = 0.7
val_size = 0.15
test_size = 0.15

# 拆分训练集与剩余集
train_ids, remainder_ids = train_test_split(
    cust_strat_key.index,
    test_size=val_size + test_size,
    stratify=cust_strat_key.values,
    random_state=42
)

# 拆分验证集与测试集
val_ids, test_ids = train_test_split(
    remainder_ids,
    test_size=test_size/(val_size+test_size),
    stratify=cust_strat_key.loc[remainder_ids].values,
    random_state=42
)

# 过滤得到最终数据集
train_df = df[df["customer_id"].isin(train_ids)]
val_df = df[df["customer_id"].isin(val_ids)]
test_df = df[df["customer_id"].isin(test_ids)]

步骤3：验证分布一致性 #

统计各集合的标签分布，确认是否符合预期：

def get_label_distribution(data_df):
    dist = data_df.groupby("label")["target"].agg(
        positive_ratio=lambda x: x.sum()/len(x)
    )
    dist["negative_ratio"] = 1 - dist["positive_ratio"]
    return dist[["positive_ratio", "negative_ratio"]]

# 查看各集合分布
print("训练集标签分布：")
print(get_label_distribution(train_df))
print("\n验证集标签分布：")
print(get_label_distribution(val_df))
print("\n测试集标签分布：")
print(get_label_distribution(test_df))

补充说明 #

• 若标签分布极度不平衡，可将罕见的分层组合合并，避免拆分时出现子集无对应样本的情况；
• 若时间序列特征重要，需结合时间维度调整拆分逻辑（比如按时间节点划分客户），同时保持分层键的分布一致性。

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