如何基于类别与ID使用Sklearn随机划分训练集与测试集？

你的问题出在GroupShuffleSplit是全局按ID分组划分，没有保证每个fruit类别内部的ID按50:50分配，所以才会出现柠檬类所有ID都进同一集合的情况。要实现「按类别内ID均分+同ID样本同组」的需求，需要针对每个类别单独处理ID划分，再合并结果。

具体实现代码 #

import pandas as pd
import random

# 设置随机种子保证结果可复现
random.seed(0)

# 生成示例数据
df = pd.DataFrame({'fruit': ['watermelon', 'watermelon', 'watermelon', 'watermelon', 'watermelon', 
'apple', 'apple', 'apple', 'apple', 'apple', 'apple', 'apple', 
"lemon", "lemon"], 
'ID': [1, 1, 1, 2, 2, 3, 4, 4, 5, 6, 6, 6 , 7 ,8], 
'value1': random.sample(range(10, 100), 14), 
'value2': random.sample(range(10, 100), 14) })

# 初始化训练/测试ID集合
train_ids = set()
test_ids = set()

# 按fruit分组，逐个处理每个类别的ID划分
for fruit, group in df.groupby('fruit'):
    # 获取当前类别下的唯一ID列表
    unique_ids = group['ID'].unique().tolist()
    # 打乱ID顺序（保证随机性）
    random.shuffle(unique_ids)
    # 计算拆分点：取ID数量的一半（整数除法自动处理奇数情况）
    split_point = len(unique_ids) // 2
    # 分配ID到训练/测试集合
    train_ids.update(unique_ids[:split_point])
    test_ids.update(unique_ids[split_point:])

# 根据ID筛选训练集和测试集
train_df = df[df['ID'].isin(train_ids)]
test_df = df[df['ID'].isin(test_ids)]

验证划分效果 #

可以运行以下代码检查结果：

# 查看训练集各类别ID分布
print("训练集各类别ID：")
for fruit, group in train_df.groupby('fruit'):
    print(f"{fruit}: {sorted(group['ID'].unique())}")

# 查看测试集各类别ID分布
print("\n测试集各类别ID：")
for fruit, group in test_df.groupby('fruit'):
    print(f"{fruit}: {sorted(group['ID'].unique())}")

# 查看各类别样本占比
print("\n训练集样本数：")
print(train_df['fruit'].value_counts())
print("\n测试集样本数：")
print(test_df['fruit'].value_counts())

关键逻辑说明 #

• 先按fruit分组，确保每个类别内部单独处理ID拆分，避免全局划分导致的类别失衡
• 对每个类别的唯一ID随机打乱后均分，保证同ID的所有样本进入同一集合
• 如果某个类别的ID数量是奇数，会自动多分配一个ID到训练集（可根据需求调整拆分逻辑，比如用round(len(unique_ids)*0.5)）

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