Получение Проверки установлено из Поезда, установленного с использованием процента от groupby () в пандах

Question

Иметь набор данных поезда с многоцелевой целевой переменной category

train.groupby('category').size()

0     2220
1     4060
2      760
3     1480
4      220
5      440
6    23120
7     1960
8    64840

Я хотел бы получить новый набор данных проверки из набора поездов, имея процент от каждого класса (скажем, 20%) во избежание пропуска классов в проверочном наборе и порчи модели.Таким образом, в основном желательным результатом будет df с такой же структурой и информацией, как у поезда, но с такими параметрами:

0     444
1     812
2     152
3     296
4      44
5      88
6    4624
7     392
8   12968

Есть ли какой-нибудь прямой подход для решения этой проблемы в пандах?

Answer 1

Groupby и семпл должны сделать это за вас

df = pd.DataFrame({'category': np.random.choice(['a', 'b', 'c', 'd', 'e'], 100), 'val': np.random.randn(100)})

idx = df.groupby('category').apply(lambda x: x.sample(frac=0.2, random_state = 0)).index.get_level_values(1)

test = df.iloc[idx, :].reset_index(drop = True)
train = df.drop(idx).reset_index(drop = True)

Редактировать: вы также можете использовать scikit learn,

df = pd.DataFrame({'category': np.random.choice(['a', 'b', 'c', 'd', 'e'], 100), 'val': np.random.randn(100)})

X = df.iloc[:, :1].values
y = df.iloc[:, -1].values
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, stratify = X)

X_train.shape, X_test.shape, y_train.shape, y_test.shape

((80, 1), (20, 1), (80,), (20,))