Converting (0,1,0, 0, 1, 1, 1) to (0, 0, 0, 1, 0, 1, 2) in R

Question

I think it is a pretty basic problem, but I would like to see if anybody can come up with a more elegant solution. Perhaps by avoiding the for loop.

I would like to have a function which takes as input a vector of 1's and 0's and returns a vector of the same length, which counts how many previous positions the same number has been.

A pretty inelegant way of doing this is:

count_me <- function(x) {
  count_vector <- numeric(length(x))
  for(i in 2:length(x)) {
    if(x[i] == x[i-1]) count_vector[i] <- count_vector[i-1] + 1
  }
  count_vector
}

which returns exactly what I want:

> (p <- sample(c(0,1), size = 10, replace = TRUE))
 [1] 0 1 0 1 1 0 0 1 1 1
> count_me(p)
 [1] 0 0 0 0 1 0 1 0 1 2

Answer 1

I believe this should work:

sequence(rle(x)$lengths) - 1