Разделить поезд и тестовый набор по категориальным столбцам

Question

У меня есть фрейм данных, содержащий около 25000 строк и 32 столбца. Я хотел бы разделить этот набор данных на поезд и тестовый тест (80/20). Однако есть некоторые столбцы с горячим кодированием. Теперь при разделении данных я хотел бы получить одинаковую долю каждого столбца с горячим кодированием в обучающем наборе.

col_1     col_2   ..  col_31    col_32
  1          0         0         0
  1          0         0         0
...
  0          0         1         0
  0          0         1         0

Итак, в обучающем наборе должно быть 80% строк, в которых каждый столбец равно 1. Я изучил различные методы разделения из Sci-kit learn, но не смог найти тот, который мог бы удовлетворить мои потребности. Есть ли у кого-нибудь решение проблемы или способный мне помочь?

Answer 1

Инициализация pandas кадр

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame(np.random.randint(0,2, size=(100, 32)))

установка 1-hot-encoded = True / False, если количество (количество нуля в строке) == 1

df['1-hot-encoded'] =  df.apply(lambda row: True if np.count_nonzero(row) == 1 else False, axis=1)

разделение с сохранением коэффициента 1-горячего кодирования

train, test = train_test_split(df, test_size=0.2, stratify=df['1-hot-encoded'])