不平衡数据集下按配送员与Target类分层划分训练测试集咨询

针对你遇到的问题——普通按Target分层无法保证每个配送员的不同类别样本同时出现在训练/测试集，核心思路是以「配送员ID + Target类别」的组合作为分层依据，这样能强制每个配送员的每个Target类别都被拆分到两个集合中（只要该组合的样本数≥2）。

方法一：基础组合分层拆分 #

直接创建「配送员+Target」的组合键，用这个键做分层拆分，确保每个组合的样本都被分配到训练和测试集：

import pandas as pd
from sklearn.model_selection import train_test_split

# 假设你的数据集存储在df中
df['stratification_key'] = df['Delivery Guy'] + '_' + df['Target'].astype(str)

# 执行分层拆分，test_size可根据需求调整
train_set, test_set = train_test_split(
    df,
    test_size=0.2,
    stratify=df['stratification_key'],
    random_state=42  # 固定随机种子保证结果可复现
)

效果说明 #

• 对于James的2个Target=1样本，拆分逻辑会自动将1个分到训练集、1个分到测试集；
• 同时全局的Target类别分布也会得到保留，适配你的不平衡数据集场景；
• 每个配送员的所有Target类别（0和1）都会同时出现在训练和测试集，让模型学习到每个配送员的完整行为模式。

方法二：处理样本数为1的特殊组合 #

如果某些配送员的某个Target类别只有1个样本（无法拆分），可以将这类样本全部归入训练集，避免测试集缺失该类别信息：

# 第一步：计算每个「配送员+Target」组合的样本量
group_counts = df.groupby(['Delivery Guy', 'Target']).size().reset_index(name='sample_count')

# 第二步：拆分出可分层的组合（样本量≥2）和不可拆分的组合（样本量=1）
valid_groups = group_counts[group_counts['sample_count'] >= 2]
valid_df = df.merge(valid_groups, on=['Delivery Guy', 'Target'], how='inner')
remaining_df = df[~df.index.isin(valid_df.index)]

# 第三步：对可分层部分执行组合拆分
train_valid, test_valid = train_test_split(
    valid_df,
    test_size=0.2,
    stratify=valid_df['stratification_key'],
    random_state=42
)

# 第四步：合并训练集（可分层训练部分 + 不可拆分样本）
train_set = pd.concat([train_valid, remaining_df], axis=0)
test_set = test_valid

验证拆分效果 #

拆分后可以用以下代码确认每个配送员的所有Target类别都同时存在于训练和测试集：

# 提取训练集和测试集的「配送员+Target」组合
train_combinations = train_set.groupby(['Delivery Guy', 'Target']).size().reset_index()
test_combinations = test_set.groupby(['Delivery Guy', 'Target']).size().reset_index()

# 检查是否存在训练集有但测试集没有的组合
missing_combinations = train_combinations[
    ~train_combinations.set_index(['Delivery Guy', 'Target']).index.isin(
        test_combinations.set_index(['Delivery Guy', 'Target']).index
    )
]

print("训练集存在但测试集缺失的配送员-Target组合：")
print(missing_combinations)

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