如何在Python中按指定类别比例划分训练集与测试集？

要让测试集中类别0占10%、类别1占90%，无法直接用train_test_split的默认分层逻辑，需要手动按类别拆分后再组合，具体步骤如下：

1. 分离不同类别的样本 #

首先把数据集按标签拆分为类别0和类别1的子集，假设你的数据是Pandas DataFrame格式（标签列名为label）：

import pandas as pd

# 假设df是包含特征和标签的完整数据集
df_0 = df[df['label'] == 0]
df_1 = df[df['label'] == 1]

2. 确定测试集各类别的样本数量 #

先设定测试集的总规模（比如你希望测试集占总数据的20%，或者直接指定样本数），再按10%/90%的比例分配两类样本数量：

# 示例：设定测试集总共有1000个样本
test_total = 1000
test_size_0 = int(test_total * 0.1)  # 测试集中类别0的数量
test_size_1 = int(test_total * 0.9)  # 测试集中类别1的数量

3. 抽取测试集和训练集样本 #

从两类子集中分别随机抽取对应数量的测试样本，剩下的作为训练样本：

# 随机抽取测试样本（random_state保证可复现）
test_0 = df_0.sample(n=test_size_0, random_state=42)
test_1 = df_1.sample(n=test_size_1, random_state=42)

# 剩余样本作为训练集
train_0 = df_0.drop(test_0.index)
train_1 = df_1.drop(test_1.index)

4. 合并并打乱数据集 #

把训练集和测试集的两类样本合并，再打乱顺序避免类别聚集：

# 合并并打乱训练集
train_df = pd.concat([train_0, train_1]).sample(frac=1, random_state=42)
# 合并并打乱测试集
test_df = pd.concat([test_0, test_1]).sample(frac=1, random_state=42)

# 拆分特征和标签
X_train, y_train = train_df.drop('label', axis=1), train_df['label']
X_test, y_test = test_df.drop('label', axis=1), test_df['label']

注意事项 #

• 如果原数据中类别0的总数量小于test_size_0（或类别1数量小于test_size_1），需要调整测试集总规模或比例，避免报错。比如可以取类别0的全部样本作为测试集，再按比例计算类别1的测试样本数：

  # 当类别0样本不足时，取全部0作为测试集，再计算1的数量
  test_0 = df_0
  test_size_1 = int(len(test_0) * 9)  # 保证0占10%，1是0的9倍
  test_1 = df_1.sample(n=test_size_1, random_state=42)
  

• 如果使用Numpy数组而非Pandas DataFrame，逻辑类似：通过布尔索引分离两类数据，再用np.random.choice抽取索引。

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