r - 在 R 中的测试和训练集中拆分数据帧时如何保留所有级别的分类变量

Question

有时，将具有分类列的数据框拆分为测试集和训练集时，训练集不会包含分类变量的所有级别。当您随后训练模型并尝试预测测试集时，预测将失败并显示:

例如:

x <- data.frame(...) # data frame with columns with very dispersed categorical variables
set.seed(123)
smp_size <- floor(0.75 * nrow(x))
train_idx <- sample(seq_len(nrow(x)), size = smp_size)
train_set <- x[train_idx, ]
test_set <- x[-train_idx, ]
m <- lm(some_formula, data=train_set)
predict(m, newdata=test_set)

Error in model.frame.default(Terms, newdata, na.action = na.action, xlev = object$xlevels) :  
    factor xxxx has new levels yyy ...

有谁知道将训练集和测试集中所有分类变量的水平设置为原始数据集中水平的简便方法？

谢谢。

Answer 1

插入符号函数 createDataPartition() 尝试处理您描述的问题。

鉴于您上面的示例，您应该能够以这种方式使用它:

train_idx <- createDataPartition(y, times = 1, p = 0.75, list=F)

这是 R 文档中关于函数 createDataPartition 的一部分:“当 y 是试图平衡拆分内的类分布的一个因素时，随机抽样在 y 的水平内完成。”