如何平衡拆分数据集保证分类变量在训练测试集均存在及R实现

我们在测试包含分类变量的插补模型准确率时，如果直接随机拆分训练集和测试集，很容易出现测试集存在训练集没有覆盖的分类变量取值的情况，导致建模时无法估计对应系数，无法正常验证模型效果。这种场景下需要在拆分时保证所有分类变量的全部取值在训练集和测试集均有分布，R中有非常简便的实现方法。

模拟测试数据 #

以下方模拟数据为例，拆分后两个数据集需要包含完全相同的字母取值集合，才能正常在测试集上验证模型准确率：

chars <- c("A","B","C","D")
complete_data <- data.frame(
  v1 = rnorm(100, 2, 100), 
  v2 = rnorm(100, 1, 100), 
  v3 = sample(chars, 100, replace = TRUE)
)

实际业务场景中数据集通常更复杂，部分分类变量的取值出现频率极低，普通随机拆分漏类别概率更高。

实现代码 #

使用caret包的createDataPartition函数做分层抽样即可完美满足需求，该函数会按指定分类变量的取值分层抽样，保证所有类别都会同时出现在训练集和测试集中：

# 加载caret包
library(caret)
# 按分类变量v3分层，抽取70%数据作为训练集
train.index <- createDataPartition(complete_data$v3, p = .7, list = FALSE)
# 生成训练集、测试集
train <- complete_data[train.index, ]
test  <- complete_data[-train.index, ]

如果需要同时满足多个分类变量都保留全取值，只需要将多个分类变量的取值拼接为一个组合分层变量，再用该变量作为分层依据传入createDataPartition即可。

内容的提问来源于stack exchange，提问作者sjillani