训练/测试拆分确保列唯一值不会在训练和测试集之间拆分？

Question

如何进行训练/测试拆分，以确保列唯一值不会在训练和测试集之间拆分？

在进行本地火车/测试拆分时，我必须按国家/地区拆分数据，以便某个国家/地区的所有年份数据都出现在火车或测试集中，但不会在两者之间拆分。

测试集中的国家代码与火车上的国家代码不同。换句话说，测试集中没有出现在火车集中的国家。因此，特定于国家/地区的功能（即国家/地区虚拟变量）将不是一种选择。但是，测试集中的国家/地区仍然与火车上的国家/地区具有相似的模式，因此其他要素工程的工作方式将与往常一样。

这是数据集的一个示例：

Country     ColumnA    ColumnB   ColumnC   Label
AB            0.2        0.5       0.1       14  
CD            0.9        0.2       0.6       60
EF            0.4        0.3       0.8       5
FG            0.6        0.9       0.2       15 

我有92个不同的国家和1400行。

我尝试了以下代码，但是稍后尝试运行简单回归时，它将返回错误。

# Train/test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=df.Country)

Answer 1

从@ArtemTurnov的评论中，我找到了答案。谢谢。

我根据国家/地区的类别拆分火车/测试数据，然后从火车和测试数据中删除国家/地区列