如何将数据集一分为二，同时确保它们在具有分类变量的列中保持平衡？

Question

我正在使用训练和测试数据集测试插补模型的准确性。我正在运行的模型使用分类变量。不幸的是，当我随机拆分数据集并在训练集上运行模型时，我无法估计测试数据集中存在的某些分类变量的系数。我想拆分数据，同时确保所有分类变量都存在于训练和测试数据集中。在 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))

在我的数据集中，问题有点棘手，因为一些分类变量非常稀缺。

编辑：

感谢您的回复。我最终按照 Antimon 的建议查找了分层抽样，并发现了 caret 包，它显然也有效。

library(caret)
train.index <- createDataPartition(complete_data$v3, p = .7, list = FALSE)
train <- complete_data[ train.index,]
test  <- complete_data[-train.index,]

Answer 1

这可以很简单地实现。

library(tidyverse)

chars<-c("A","B","C","D")

complete_data <- tibble(v1=rnorm(100,2,100), 
                        v2=rnorm(100,1,100), 
                        v3=sample(chars, 100, replace=TRUE))


propCategory = function(data, category, prop){
  category = enquo(category)
  cat1 = data %>% pull(!!category)
  unlist(sapply(as.list(unique(cat1)), function(x) {sample(which(cat1==x), sum(cat1==x)*prop)}))
}
complete_data %>% propCategory(v3, .2)

输出

 [1]  98  35  20  78  40  70  87   3  86  38  22 100  80  93  47   5  24  29  26

如您所见，我的propCategory 函数返回轴向索引。但是让我们检查一下它们是否包含您需要的东西。 首先，让我们检查一下训练指标。

train = complete_data %>% propCategory(v3, .75)

complete_data[train,] %>% distinct(v3)
complete_data[train,] %>% nrow()

输出

> complete_data[train,] %>% distinct(v3)
# A tibble: 4 x 1
  v3   
  <chr>
1 B    
2 A    
3 D    
4 C    
> complete_data[train,] %>% nrow()
[1] 74

现在是测试索引的时候了。

complete_data[-train,] %>% distinct(v3)
complete_data[-train,] %>% nrow()

输出

> complete_data[-train,] %>% distinct(v3)
# A tibble: 4 x 1
  v3   
  <chr>
1 B    
2 A    
3 D    
4 C    
> complete_data[-train,] %>% nrow()
[1] 26

如您所见，训练和测试数据都包含您的每个类别。

关于prop 参数的一点说明。 我的 propCategory 函数的编写方式是，对于变量类别中的每个值，它返回带有 prop * 的随机选择索引的数量（分类变量的保存值的数量）。

仔细看看下面的结果。

complete_data %>% group_by(v3) %>% 
  summarise(n = n(), prop = n()/nrow(.))

complete_data[train,] %>% group_by(v3) %>% 
  summarise(n = n(), prop = n()/nrow(.))

complete_data[-train,] %>% group_by(v3) %>% 
  summarise(n = n(), prop = n()/nrow(.))

输出

> complete_data %>% group_by(v3) %>% 
+   summarise(n = n(), prop = n()/nrow(.))
# A tibble: 4 x 3
  v3        n  prop
  <chr> <int> <dbl>
1 A        26  0.26
2 B        35  0.35
3 C        24  0.24
4 D        15  0.15
> complete_data[train,] %>% group_by(v3) %>% 
+   summarise(n = n(), prop = n()/nrow(.))
# A tibble: 4 x 3
  v3        n  prop
  <chr> <int> <dbl>
1 A        19 0.257
2 B        26 0.351
3 C        18 0.243
4 D        11 0.149
> complete_data[-train,] %>% group_by(v3) %>% 
+   summarise(n = n(), prop = n()/nrow(.))
# A tibble: 4 x 3
  v3        n  prop
  <chr> <int> <dbl>
1 A         7 0.269
2 B         9 0.346
3 C         6 0.231
4 D         4 0.154

Answer 2

有几种方法可以做到这一点。您需要将您的数据除以v3，然后随机拆分每个组：

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))

现在我们将使用by() 函数将数据按v3 分组，并随机抽取每组中一半的行名样本：

test <- as.numeric(unlist(by(complete_data, complete_data$v3, function(x) sample(rownames(x), length(rownames(x))/2))))
train_test <- rep("train", nrow(complete_data))
train_test[test] <- "test"
table(complete_data$v3, train_test)
#    train_test
#     test train
#   A   11    12
#   B   12    13
#   C   13    14
#   D   12    13

现在complete_data[train_test=="train", ] 是您的训练集，complete_data[train_test=="test", ] 是您的测试集。