根据分类列拆分训练和测试集

Posted 2023-01-14

我有一个包含大约25000行和32列的数据框。我想将此数据集分为训练和测试测试（80/20）。但是，某些列是1热编码的。现在，当分割数据时，我希望将每个1-hot编码列的相同比例放入训练集中。

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学习中了解了不同的拆分方法，但是找不到能够满足我需要的拆分方法。有没有人提供解决方案或者可以帮助我的人？​​

答案

初始化大熊猫框

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热编码=真/假

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'])

希望这会有所帮助