如何按类别标签手动拆分数据集为指定比例的训练集与测试集

你要的是按类别分层拆分数据集的效果，有两种常用实现方式：

方案1：直接用train_test_split的分层参数 #

你之前用train_test_split没实现需求是因为没开启分层参数，只要传入stratify参数指定标签列，就能自动保证训练集、测试集里每个类别的比例和原数据集完全一致：

from sklearn.model_selection import train_test_split
# 假设X是特征矩阵，y是对应的类别标签，test_size是测试集占比
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, stratify=y, random_state=42)

注：random_state参数用来固定随机种子，保证拆分结果可复现。

方案2：手动按类别拆分（自定义程度更高） #

如果需要对每个类别的拆分逻辑做自定义调整，可以遍历每个类别单独拆分后再合并：

import pandas as pd
from sklearn.model_selection import train_test_split

# 假设数据集A是DataFrame格式，label_col为类别列的列名
def split_by_class(dataset, label_col, test_size=0.5):
    train_dfs = []
    test_dfs = []
    # 遍历每个类别单独执行拆分
    for label in dataset[label_col].unique():
        class_df = dataset[dataset[label_col] == label]
        class_train, class_test = train_test_split(class_df, test_size=test_size, random_state=42)
        train_dfs.append(class_train)
        test_dfs.append(class_test)
    # 合并所有类别的训练集、测试集
    train_df = pd.concat(train_dfs).reset_index(drop=True)
    test_df = pd.concat(test_dfs).reset_index(drop=True)
    return train_df, test_df

# 调用示例：按50%比例拆分训练集和测试集
train_set, test_set = split_by_class(A, label_col="你的类别列名", test_size=0.5)

普通场景用方案1足够满足需求，如果有特殊的类别拆分规则可以基于方案2调整逻辑。

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