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

【问题标题】：Split train and test set according to categorical columns根据分类列拆分训练和测试集
【发布时间】：2020-09-12 01:22:46
【问题描述】：

我有一个包含大约 25000 行和 32 列的数据框。我想将此数据集拆分为训练和测试测试（80/20）。但是，有些列是 1-hot 编码的。现在，在拆分数据时，我希望将每个 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 learn 中查看了不同的拆分方法，但无法找到一种可以满足我需求的方法。有没有人有解决方案或能够帮助我？

【问题讨论】：

据我了解，可以使用sklearn的train_test_split函数中的stratify参数来获取分层样本（比例）。
@dheinz 我不确定如何在选定的几列上使用stratify
检查这是否有帮助 - [链接] (*.com/questions/45516424/…)

标签： python pandas scikit-learn

【解决方案1】：

初始化熊猫框架

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 if count(no.of.zero in a row) == 1

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

在保持 1-hot-encoded 比率的同时进行拆分

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

【讨论】：

这会将完整的df['1-hot-encoded'] 设置为False
1) 检查 ur 列的数据类型是否为整数
2) 如果您在 pd.df 中使用随机值，可能是您没有有效的 1-hot-encoded 行跨度>