Lagging vector adding 1 while resetting to 0 when a condition is met

Question

I have a sequence of treatments, one per day (binary), say:

trt <- c(0, 0, 1, 0, 0, 0, 1, 0, 0)

I want to create a vector, days_since, that:

1. Is NA up until the first treatment.
2. Is 0 where trt is 1
3. Counts the days since the last treatment

So, the output days_since should be:

days_since <- c(NA, NA, 0, 1, 2, 3, 0, 1, 2)

How would I do this in R? To get days_since, I basically need to lag by one element and add 1, but resetting every time the original vector (trt) is 1. If this is doable without a for-loop, that would be ideal, but not absolutely necessary.

Answer 1

Maybe you can try the code below

v <- cumsum(trt)
replace(ave(trt,v,FUN = seq_along)-1,v<1,NA)

which gives

[1] NA NA  0  1  2  3  0  1  2

---

Explanation

• First, we apply cumsum over trt to group the treatments

> v <- cumsum(trt)
> v
[1] 0 0 1 1 1 1 2 2 2

• Secondly, using ave helps to add sequential indices within each group

> ave(trt,v,FUN = seq_along)-1
[1] 0 1 0 1 2 3 0 1 2

• Finally, since the value is NA before the first treatment, it means all the value before v == 1 appears should be replaced by NA. Thus we use replace, and the index logic follows v < 1

> replace(ave(trt,v,FUN = seq_along)-1,v<1,NA)
[1] NA NA  0  1  2  3  0  1  2