Pandas 根据目标变量及其簇分层拆分为训练集、测试集和验证集

Question

我有一个包含一些特征的数据框和一个属于 {0,1} 的目标列。我需要将这个数据集拆分为训练集、测试集和验证集。验证部分必须是数据集的 20%，其余 80% 必须拆分，以便其中的 80% 进入训练集。这可以通过 sklearn 的 train_test_split

轻松实现

我的问题是拆分必须以分层的方式进行 基于集群 我计算了两个目标值。

为了计算集群，我将两个目标的条目分成两个子集，例如

ones = df[df_numerical['Target'] == 1].copy()
zeroes = df[df_numerical['Target'] == 1].copy()

然后对于每个子集，我使用 kmeans 来计算它们的集群，并将集群添加到数据框中，例如：

# the number of clusters for both variables is not the same
clusters_1 = kmeans_1.predict(ones[NUMERICAL_FEATURES])
ones['Cluster'] = clusters_1

clusters_0 = kmeans_0.predict(zeroes[NUMERICAL_FEATURES])
zeroes['Cluster'] = clusters_0

现在我如何拆分数据集，以便它们按簇大小分层？

我需要的拆分必须以这种方式完成：假设有 100 条记录，80 条 class 1 和 20 条 class 0，我需要将这些记录拆分为 70 / 30%，所以我需要 class 1 的 56（80% 的 70%）记录和 class0 的 14（20% 的 70%）记录。我知道这可以使用 [= train_test_split 的 16=] 参数，但我的问题是除此之外，拆分还必须分层 w.r.t 每个目标值的聚类。

我认为的一个解决方案是提取两个 classes 的元素索引，将它们放入列表中，从中提取正确数量的元素，然后重新组合数据帧：

cluster_indices_0 = zeroes.groupby(['Cluster']).apply(lambda x: x.index)
cluster_indices_1 = ones.groupby(['Cluster']).apply(lambda x: x.index)

但是通过这种方式我必须手动计算每个集群要弹出的元素数量，我正在寻找一种自动执行此操作的方法。

sklearn 或 pandas 中是否有一个函数可以在计算要提取的元素数量时无需获取列表来实现我正在寻找的功能？

Answer 1

由于您的数据已经按目标拆分，您只需在每个子集上调用 train_test_split 并使用聚类列进行分层。

train_test_0, validation_0 = train_test_split(zeroes, train_size=0.8, stratify=zeroes['Cluster'])
train_0, test_0 = train_test_split(train_test_0, train_size=0.7, stratify=train_test_0['Cluster'])

然后对目标 1 执行相同的操作并组合所有子集

Pandas 根据目标变量及其簇分层拆分为训练集、测试集和验证集

Pandas stratified splitting into train, test, and validation set based on the target variable its cluster

python

sampling

dataframe

pandas