如何按照某些标准将数据集拆分为子集？

Question

虽然我使用与机器学习相关的术语，但我的问题是 100% 工程主题，与统计和数学无关。因此，我在这个论坛而不是 Cross Validated 中询问它。

这是我将用来评论我的问题的示例数据：

X = pd.DataFrame(columns=["F1","F2"], 
                  data=[[23,0.8],
                        [11,5.35],
                        [24,19.18],
                        [15,10.25],
                        [10,11.30],
                        [55,44.85],
                        [15,33.88],
                        [12,45.30],
                        [14,22.20],
                        [15,15.80],
                        [83,0.8],
                        [51,5.35],
                        [34,30.28],
                        [35,15.25],
                        [60,13.30],
                        [75,44.80],
                        [35,30.77],
                        [62,40.33],
                        [64,23.40],
                        [14,11.80]])

y = pd.DataFrame(columns=["y"], 
                  data=[[0],
                        [0],
                        [1],
                        [0],
                        [2],
                        [2],
                        [2],
                        [1],
                        [0],
                        [1],
                        [0],
                        [0],
                        [1],
                        [0],
                        [1],
                        [0],
                        [1],
                        [1],
                        [0],
                        [2]])

我应该将数据分成训练集和测试集。一个经典的方法是使用sklearn的train_test_split函数：

X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y, test_size=0.25)

但我想指定要分配给训练集和测试集的记录百分比。更多细节解释如下。

在我的例子中，我处理了一个多类分类问题，其中y 可能取 3 个不同的值之一：0、1、2。值为 2 的记录非常罕见（在我的真实数据集中，约占整个数据集的 3%）。因此这是一个不平衡的分类问题。

由于这是一个不平衡的分类问题，稀有类的记录非常重要。因此，我想更新model_selection.train_test_split 如下：我想为训练集和测试集分配每个班级的记录百分比。 例如， 表示将稀有类的 90% 的记录分配给训练集。

在我的示例中，例如，我想在训练集（X_train、y_train）中获取 3 条 y 等于 2 的记录，并在测试集中获取 1 条记录。

我用谷歌搜索了类似的问题，但没有找到任何东西。

为了解决这个任务，我打乱了初始数据框：

df = pd.concat([X, y], axis=1)

df = df.sample(frac=1).reset_index(drop=True)

但是，我不知道如何继续其他任务。也许有一些 sklearn 内置函数或一些库可以解决这个问题？

Answer 1

在 train_test_split 中有一个称为分层的选项。也看看这个kind of similar question

要完成您需要的比例，您可以使用 numpy 中的np.random.choice：

import numpy as np
df = pd.concat([X,y], axis = 1)

#get index values for y = 0
i0 = np.random.choice(df.loc[df.y==0].index.values,
round(len(df.loc[df.y==0])*.5), replace = False)

i1 = np.random.choice(df.loc[df.y==1].index.values,
round(len(df.loc[df.y==1])*.6), replace = False)

i2 = np.random.choice(df.loc[df.y==2].index.values,
round(len(df.loc[df.y==1])*.9), replace = False)

df_train = df.loc[df.index.isin(np.concatenate([i1,i2,i0]))]
df_test = df.loc[~df.index.isin(np.concatenate([i1,i2,i0]))]