Pandas，scikit - 拆分火车并进行条件测试答案

【问题标题】：Pandas, scikit - split train and test with conditionPandas，scikit - 拆分火车并进行条件测试
【发布时间】：2020-05-11 07:57:20
【问题描述】：

我有一个数据框，data，具有以下结构（我的实际上要大得多，但这只是为了说明目的）：

我正在使用 scikit-learn 来拆分数据：

train, test = train_test_split(data, test_size=test_size)

但是，我想找到一种方法来拆分数据，以保证我在两组中每个标签至少有一行。对于示例数据集，这意味着有这样的东西（当然是打乱的）：

火车
a b c tag A 3 2 4 B 2 1 3 A 4 3 2 B 4 1 1

测试 similar but with the remaining elements (according to the proportion)

基本上，我希望在两组中都有整个范围/种类的标签。

我提前感谢所有帮助。

【问题讨论】：

检查函数内部参数stratify =

标签： python pandas scikit-learn

【解决方案1】：

我想这就是你所追求的。

import numpy as np
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

有关此主题的更多信息，请参阅下面的链接。

https://towardsdatascience.com/train-test-split-and-cross-validation-in-python-80b61beca4b6

【讨论】：

这个答案似乎并不像 OP 所要求的那样“保证我在两组中的每个标签中至少有一行”。