Pandas, scikit - разделенный поезд и тест с условием

Question

У меня есть датафрейм, данные со следующей структурой (моя на самом деле намного больше, но это только для иллюстрации):

a b c tag
A 3 2 4
B 2 1 3
A 5 3 3
A 4 3 2
B 2 4 3
A 3 5 2
B 4 1 1
C 2 3 1
C 1 3 4
B 5 2 4

Я использую scikit-learn для разделения данных :

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

Однако я хочу найти способ разбить данные таким образом, чтобы гарантировать, что у меня есть хотя бы одна строка каждого тега в обоих наборах. Для примера набора данных это означало бы наличие чего-то подобного (но, конечно, перемешанного):

train
a b c tag A 3 2 4 B 2 1 3 A 4 3 2 B 4 1 1

test similar but with the remaining elements (according to the proportion)

В основном Я хотел бы иметь весь диапазон / разнообразие тегов в обоих наборах.

Заранее благодарен за помощь.

`

Answer 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)

См. Ссылку ниже для получения дополнительной информации об этом топи c.

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