如何将数据按组拆分？

Question

我有以下数据

user_id, item_id
1, 2
1, 3
2, 5,
2, 4
1, 4

我想用一个约束来分割数据：

每个用户都必须在火车组或测试组中，但不能同时处于这两种状态。

所以[[1,2], [1,3], [1,4]] / [[2,5],[2,4]]是有效的拆分

但是[[1,2],[2,5][1,4]] / [[1,3]... ]不是有效的拆分，因为用户已经被拆分了。

我不确定from sklearn.model_selection import train_test_split是否支持它。

Answer 1

如果你想要有固定比例的用户在培训和测试，你需要得到独特的用户和分裂的用户。以下伪码

train_ids, test_ids = train_test_split(unique_users)
train_X = data[data.user_id==train_ids]
test_X = data[data.user_id==test_ids]