python - 如何按组拆分数据？

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支持它..

标签： pythonscikit-learn

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]