Инициализация 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'])