拆分火车测试数据集，将相似的值保持在一起

问题描述

我有一个带有ID的动物类型的数据集，我想将所述数据集分解为“测试/训练”数据集。我还希望将相应动物的所有ID保留在Train或Test数据集中。下面是一个数据示例，火车/测试的随机分配比率为80/20。

Animal  ID  Test/Train
CAT 1   TRAIN
CAT 1   TRAIN
CAT 2   TRAIN
CAT 2   TRAIN
CAT 3   TRAIN
CAT 3   TEST
CAT 4   TRAIN
CAT 4   TRAIN
CAT 5   TEST
CAT 5   TRAIN
DOG 1   TRAIN
DOG 1   TRAIN
DOG 2   TRAIN
DOG 2   TRAIN
DOG 3   TRAIN
DOG 3   TRAIN
DOG 4   TEST
DOG 4   TEST
DOG 5   TRAIN
DOG 5   TRAIN

请注意在训练和测试数据集中如何存在ID为3和ID 5的CAT。 scikit-learn train_test_split中是否有一个功能，可以在保持测试比率的同时，将所有相似值保留在同一训练/测试数据集中的列中？因此，如果ID为3的CAT将一个值标记为Train数据，那么其他所有CAT和ID 3的其他记录也将标记为Train data。

解决方法

是否将stratify参数保留为yes，然后删除并检查。