如何根据一些准则将数据集划分为子集

Question

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

这是我的样本数据，我将用它来评论我的问题：

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如下：我想为训练集和测试集分配每个类%的记录例如，<50%、60%、90%>意味着稀有类90%的记录被分配到训练集中

在我的示例中，例如，我希望在训练集中获得y等于2的3条记录(X_train，y_train)，在测试集中获得1条记录。

我在谷歌上搜索了类似的问题，但什么也没找到。

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

df = pd.concat([X, y], axis=1)
df = df.sample(frac=1).reset_index(drop=True)

然而，我不知道如何处理剩下的任务。也许有一些sklearn内置函数或库可以解决这个问题？

Answer 1

在train_testrongplit中有一个名为分层的选项。也看看这类类似的问题

要完成所需的比例，可以使用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]))]