R语言中将向量随机拆分为k个块的自定义实现可优化吗？

问题背景 #

我见过很多「在R中将向量X拆分为Y个块」相关问题的变体，当我需要将向量拆分为Y个大小随机的块时，发现这类带随机性要求的拆分需求没有现成的实现方案，因此自行编写了如下函数：

k.chunks = function(seq.size, n.chunks) {
  break.pts = sample(1:seq.size, n.chunks, replace=F) %>% sort() # 从向量长度范围内无放回采样断点，避免重复选择
  groups = rep(NA, seq.size) # 初始化空输出向量
  groups[1:break.pts[1]] = 1 # 第一组起点固定为1，单独赋值

for (i in 2:(n.chunks)) { # 遍历其余分组
    groups[break.pts[i-1]:break.pts[i]] = i # 赋值对应分组编号
    }
    groups[break.pts[n.chunks]:seq.size] = n.chunks # 最后一组终点固定为向量长度，单独赋值
    return(groups)
    }

我的疑问如下：该实现是否存在不够简洁或效率低下的问题？我后续代码会数千次调用该函数，因此效率对我而言十分重要，若能避免for循环、无需手动分别设置首尾分组就更好了。另外除了n.chunks大于seq.size的情况外，还有哪些逻辑合法的输入会导致该函数运行出错？

解答 #

一、原实现存在的问题 #

1. 核心逻辑缺陷：拆分N个块仅需要N-1个断点即可，原函数采样了N个断点，会导致拆分结果不符合预期。比如seq.size=2、n.chunks=2时，原函数返回值为c(2,2)，完全没有实现拆分效果。
2. 性能偏低：for循环在seq.size较大、调用次数较多时性能损耗明显，同时逐段赋值属于非向量化操作，运行速度慢。
3. 边界处理冗余：单独对首尾分组赋值属于不必要操作，可通过向量化函数统一处理。

二、优化实现方案 #

完全避免for循环，采用R内置向量化函数实现，运行效率比原函数高1~2个数量级，代码更简洁：

k.chunks_fast <- function(seq.size, n.chunks) {
  # 前置参数校验，避免非法输入
  stopifnot(seq.size >= n.chunks, n.chunks >= 1, seq.size >= 1)
  if (n.chunks == 1) return(rep(1L, seq.size))
  # 采样n.chunks-1个不重复断点
  break_pts <- sort(sample(seq_len(seq.size - 1), n.chunks - 1, replace = FALSE))
  # 直接生成全部分组编号
  findInterval(seq_len(seq.size), break_pts) + 1L
}

三、其他会导致原函数出错的输入场景 #

• n.chunks == 1：虽不会直接报错，但原函数采样1个断点，会对同一区间重复赋值1，属于冗余操作
• seq.size == 1：仅当n.chunks ==1时能正常运行，其余情况直接报错
• n.chunks == 0：采样步骤直接报错
• seq.size <= 0：初始化分组向量的步骤直接报错
• 非整数输入seq.size或n.chunks：参数被自动截断后运行结果不符合预期

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