按组将DataFrame拆分为大小相近的块并均分各组至n个拆分

原始数据与需求 #

首先定义示例数据：

library(tidyverse)
df <- data.frame(x=rep(c('A','B'), times=c(5,15)), y=1:20)

需求说明：将df拆分为n个大小相近的子数据框，需满足两个约束：

• x的每个分组必须出现在所有n个拆分结果中
• 每个x分组内的数据需均分到这n个拆分里

例如当n=3时，预期拆分结果如下：

$`1`
   x  y
1  A  1
2  A  2
3  B  6
4  B  7
5  B  8
6  B  9
7  B 10

$`2`
   x  y
1  A  3
2  A  4
3  B 11
4  B 12
5  B 13
6  B 14
7  B 15

$`3`
   x  y
1  A  5
2  B 16
3  B 17
4  B 18
5  B 19
6  B 20

其中A组被拆分为[2,2,1]的块，B组被拆分为[5,5,5]的块。

问题分析 #

你尝试的代码仅生成了两个拆分，原因是group_by(x)后直接用split(., cut(seq_along(.), n, labels = FALSE))会对每个分组单独拆分，cut的方式无法保证稳定生成n个拆分，且最终没有统一的拆分ID来合并分组结果。

解决方案 #

正确的思路是先给每个分组内的行分配拆分ID，再按ID拆分整个数据框：

n <- 3

# 生成拆分ID并拆分
split_result <- df %>%
  group_by(x) %>%
  mutate(split_id = rep(1:n, length.out = n())) %>%
  ungroup() %>%
  split(.$split_id)

# 查看结果
split_result

代码逻辑解释 #

• rep(1:n, length.out = n())：针对每个x分组，循环生成1到n的序列，长度匹配当前分组的行数。这样能保证每个分组的行被均匀分配到n个拆分中，且每个拆分都会包含该分组的行（只要分组行数≥1）。
• 取消分组后，按生成的split_id拆分数据框，就能得到符合要求的n个子数据框。

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