如何将分组数据集拆分为不超指定大小的块且同组不拆分？

需要将数据集拆分为行数不超过指定最大值的块，核心约束为同一id组的记录必须完整落在同一个块中，不能拆分到不同块。具体示例如下：

示例数据集 #

df <- tibble(id = c("id1", "id1", "id2", "id2", "id2", "id3", "id4", "id5", "id5", "id6"),
             group_size = c(2, 2, 3, 3, 3, 1, 1, 2, 2, 1))

原始数据结构：

id    group_size
   <chr>      <dbl>
 1 id1            2
 2 id1            2
 3 id2            3
 4 id2            3
 5 id2            3
 6 id3            1
 7 id4            1
 8 id5            2
 9 id5            2
10 id6            1

期望输出（块大小最大值设为3） #

新增part列作为块标识，结果如下：

id    group_size  part
   <chr>      <dbl> <dbl>
 1 id1            2     1
 2 id1            2     1
 3 id2            3     2
 4 id2            3     2
 5 id2            3     2
 6 id3            1     3
 7 id4            1     3
 8 id5            2     4
 9 id5            2     4
10 id6            1     4

请问是否有现成的R函数或包可以实现该功能？如果没有，该如何编程实现？

1. 现有工具情况 #

目前没有专门直接实现该逻辑的通用R包或函数，因为该需求属于带约束的分组拆分（组不可拆分、块大小上限），需要自定义逻辑来实现。

2. 自定义实现代码 #

核心思路：先提取每个id组的唯一大小，通过累加判断的方式给每个组分配块标识，最后将块标识映射回原数据集。

代码示例 #

library(tidyverse)

# 定义拆分函数
split_into_blocks <- function(data, max_block_size, group_col = "id", size_col = "group_size") {
  # 提取每个组的唯一大小
  group_sizes <- data %>%
    distinct({{group_col}}, {{size_col}}) %>%
    rename(group = {{group_col}}, size = {{size_col}})
  
  # 初始化块标识和当前块累计行数
  current_block <- 1
  current_sum <- 0
  group_sizes$part <- NA_integer_
  
  # 遍历每个组分配块
  for (i in seq_len(nrow(group_sizes))) {
    group_size <- group_sizes$size[i]
    # 若当前组加入后超过上限且当前块已有内容，则开启新块
    if (current_sum + group_size > max_block_size && current_sum > 0) {
      current_block <- current_block + 1
      current_sum <- group_size
    } else {
      current_sum <- current_sum + group_size
    }
    group_sizes$part[i] <- current_block
  }
  
  # 将块标识合并回原数据
  data %>%
    left_join(group_sizes, by = setNames("group", group_col)) %>%
    select(-size)
}

# 测试函数
df <- tibble(id = c("id1", "id1", "id2", "id2", "id2", "id3", "id4", "id5", "id5", "id6"),
             group_size = c(2, 2, 3, 3, 3, 1, 1, 2, 2, 1))

split_into_blocks(df, max_block_size = 3)

代码说明 #

• 用distinct获取每个id的组大小，避免重复计算组规模
• 循环中通过累加判断是否需要新建块，严格保证同一组不会被拆分到不同块
• 最后通过left_join将块标识映射回原数据集，确保所有行都能匹配到对应的part值

运行上述代码可得到与示例完全一致的输出结果。

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