我有一个如下所示的样本数据:
import pandas as pddf = pd.DataFrame({"x": [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120], "id": [1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5], "label": ["a", "a", "a", "b", "a", "b", "b", "b", "a", "b", "a", "b"]})所以我的数据看起来像这样
x id label 10 1 a 20 1 a 30 1 a 40 1 b 50 2 a 60 2 b 70 3 a 80 3 a 90 4 b100 4 a110 5 b120 5 a我想根据标签分布将这些数据分成两个组（训练集和测试集），给定测试样本的数量（例如6个样本）。我的设置更喜欢将测试集的大小定义为表示测试样本数量的整数，而不是百分比。然而，在我的特定领域中，任何id必须只分配到一个组中。例如，如果id 1被分配到训练集，其他带有id 1的样本不能被分配到测试集。因此，预期的输出是如下的两个数据框:
训练集
x id label 10 1 a 20 1 a 30 1 a 40 1 b 50 2 a 60 2 b测试集
x id label 70 3 a 80 3 a 90 4 b100 4 a110 5 b120 5 a训练集和测试集具有相同的类别分布（a:b是4:2），id 1和2只被分配到训练集，而id 3、4、5只被分配到测试集。我曾经使用sklearn的train_test_split来处理，但我不确定如何在这种条件下应用它。您能提供一些处理这种情况的建议吗？
回答：
sklearn.model_selection除了train_test_split之外还有其他几个选项。其中一个选项旨在解决您所追求的目标。在这种情况下，您可以使用GroupShuffleSplit，正如文档中提到的，它提供随机的训练/测试索引来根据第三方提供的组别分割数据。对于这种情况，您还可以使用GroupKFold，它非常有用。
from sklearn.model_selection import GroupShuffleSplitX = df.drop('label',1)y=df.label现在您可以实例化GroupShuffleSplit，并像使用train_test_split一样操作，唯一的区别是指定一个group列，该列将用于根据组别值分割X和y:
gs = GroupShuffleSplit(n_splits=2, test_size=.6, random_state=0)train_ix, test_ix = next(gs.split(X, y, groups=X.id))现在您可以索引数据框来创建训练集和测试集:
X_train = X.loc[train_ix]y_train = y.loc[train_ix]X_test = X.loc[test_ix]y_test = y.loc[test_ix]结果如下:
print(X_train) x id4 50 25 60 28 90 49 100 410 110 511 120 5而测试集为:
print(X_test) x id0 10 11 20 12 30 13 40 16 70 37 80 3