不规则分组的组内索引计算：如何用向量化替代循环提效？

给定一组不规则分组标签数据g，需要生成每个分组内的递增索引k（即重置累加，每个新分组从1开始计数），示例如下：

g = c(1,1,1, 2, 3,3, 4, 5, 6,6,6,6,6, 7, 8, 9,9,9,9, 10, 11, 12, 13,13)
k = c(1,2,3, 1, 1,2, 1, 1, 1,2,3,4,5, 1, 1, 1,2,3,4,  1,  1,  1,  1, 2)

已有可行的循环实现方案，但由于循环在处理大数据时效率较低，希望找到更高效的向量化方案：

g.index = function(g){
  rep.i = c(F,diff(g)==0)
  k = numeric(length(g))
  for (i in 1:length(g)){
    if (rep.i[i]){ cs = cs + 1 } else { cs = 1 }
    k[i] = cs
  }
  return(k)
}

高效向量化实现方法 #

方法1：rle() + sequence()（base R最快方案之一） #

利用R内置的游程编码函数rle()获取每个连续分组的长度，再用sequence()生成对应组内序列，代码极简且效率极高：

k <- sequence(rle(g)$lengths)

原理：rle(g)返回的列表中，lengths元素记录了每个连续相同值的长度（示例中为c(3,1,2,1,1,5,1,1,4,1,1,1,2)），sequence()会将这些长度转换为逐个递增的序列，正好对应每个组内的索引。

方法2：ave()函数（base R简洁方案） #

ave()是base R中专门用于分组操作的函数，直接按g分组并生成组内序列：

k <- ave(rep(1, length(g)), g, FUN = seq_along)

原理：对每个g的分组，应用seq_along生成从1开始的连续整数，自动完成分组内的计数。

方法3：cumsum()手动实现（base R底层向量化） #

通过标记分组变化点，结合累加实现重置计数：

# 标记每个新分组的起始位置（第一个元素默认是新分组）
group_starts <- c(TRUE, diff(g) != 0)
# 生成唯一分组ID，再对每个ID生成组内索引
group_ids <- cumsum(group_starts)
k <- ave(group_ids, group_ids, FUN = seq_along)

效率对比 #

上述向量化方法的执行效率远高于循环实现，尤其是当数据量达到百万级以上时，循环的开销会被显著放大，而向量化操作利用了R的底层优化，速度提升可达几十倍甚至上百倍。

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