Count vector elements in range when condition is met

Question

I have a vector of numbers:

x <- c(0, 0, 0, 30, 60, 0, 0, 0, 0, 0, 10, 0, 0, 15, 45, 0, 0)

For each element i in x, I would like to do the following

1. If x[i] > 0, return 0
2. If all 4 elements before x[i] are 0, return NA
3. If the 4 elements before x[i] are not 0, count how many elements are between the last not-0 element and x[i]

I expect this output:

#> x
#[1]  0  0  0 30 60  0  0  0  0  0 10  0  0 15 45  0  0
#> x_out
#[1] NA NA NA  0  0  1  2  3  4 NA  0  1  2  0  0  1  2

Notice that the solution should also work when there are less than 4 elements available at the beginning of the vector (i.e. condition 2 and 3 should use as many elements as are available). 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

Here is a simple Rcpp solution. Create a new C++ file in RStudio and paste the code into it and source the file. Obviously, you'll need to have installed Rtools if you use Windows.

#include <Rcpp.h>
using namespace Rcpp;    

// [[Rcpp::export]]
IntegerVector funRcpp(const IntegerVector x) {
  const double n = x.length();
  int counter = 4;
  IntegerVector y(n);

  for (double i = 0; i < n; ++i) {
    if (x(i) > 0) {
      y(i) = 0;
      counter = 0;
    }
    else {
      if (counter > 3) {
        y(i) = NA_INTEGER;
      } else {
        counter++;
        y(i) = counter;
      }
    }
  }

  return y;
}


/*** R
x <- c(0, 0, 0, 30, 60, 0, 0, 0, 0, 0, 10, 0, 0, 15, 45, 0, 0)
funRcpp(x)
*/

This returns the desired result:

> funRcpp(x)
 [1] NA NA NA  0  0  1  2  3  4 NA  0  1  2  0  0  1  2