Within-group index (irregular groups)

Question

I have some data (irregular group labels) like g, and I want to obtain k -- i.e. within-group indices, via resetting cumsum.

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)

I have a working solution:

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

But I'm worried it will be slow due to loops versus vectorization. Is there a more efficient way?

user_6762755 · Answer

As commented by @akrun, use data.table::rowid

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)

library(data.table)

all(rowid(g) == k)
#> [1] TRUE

^{Created on 2022-07-29 by the reprex package (v2.0.1)}

user_5228288 · Answer

As with many things in R, you can't trust your intuition.

Here is a one-liner using lapply and split:

g.index.2 = function(g){
  k = unlist(lapply(split(rep(1,length(g)),g),cumsum))
}

We can see they give the same results using (graph below):

plot(g.index(g))
lines(g.index.2(g))

However, the first function is actually much faster. I'm not sure why.

N = 1e4
g = sort(ceiling(runif(N)*N/10))
microbenchmark::microbenchmark({g.index(g)},{g.index.2(g)},times=1e6/N)

Unit: microseconds
                 expr      min       lq     mean   median       uq       max neval cld
 {       g.index(g) }  697.566  760.452 1078.655  907.409 1090.495  6857.132   100  a
 {     g.index.2(g) } 7722.124 8017.586 8325.830 8258.430 8563.026 10101.035   100   b

user_19651015 · Answer

You can break your problem down into the following steps that compose into another one-line solution:

Determine run lengths using rle()
Generate an integer sequence for each run length using seq_len()
Concatenate the sequences into a single integer vector using do.call(c, ...)

do.call(c, lapply(rle(g)$lengths, seq_len))
#>  [1] 1 2 3 1 1 2 1 1 1 2 3 4 5 1 1 1 2 3 4 1 1 1 1 2

Within-group index (irregular groups)

3 Answers