如何按组拆分数据？

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]