Reverse count vector elements when sequence occurs

Question

I have a vector with values NA, 0, and 1:

x <- c(NA, 0, 0, 1, 1, 1, 1, NA, 0, 0, 0, 0, NA, NA, 1, 1, 1, NA)
#> x
#[1] NA  0  0  1  1  1  1 NA  0  0  0  0 NA NA  1  1  1 NA

Whenever the sequence switches from 1 to NA, I would like to count the positions of non-NAs before that event and replace the elements with that number. I expect this output:

#> x_output
#[1] NA  6  5  4  3  2  1 NA  0  0  0  0 NA NA  3  2  1 NA

Does anybody have a solution for this? A vectorised approach is preferred because the vectors are long and the dataset is fairly big.

Answer 1

Using rle to define run lengths and ave to create the sequences:

x <- c(NA, 0, 0, 1, 1, 1, 1, NA, 0, 0, 0, 0, NA, NA, 1, 1, 1, NA)

fun <- function(x) {
  x <- rev(x)
  y <- rle(!is.na(x))
  y$values[y$values] <- seq_along(y$values[y$values])
  y <- inverse.rle(y)

  x[!is.na(x)] <- ave(x[!is.na(x)], y[!is.na(x)], FUN = function(x) {
    if (x[1] == 0L) return(x)
    seq_along(x)
  })
  rev(x)
}

fun(x)
#[1] NA  6  5  4  3  2  1 NA  0  0  0  0 NA NA  3  2  1 NA