Cumsum ignoring NA's with reset

Question

I have conditional sum with reset at zero.

criteria1 <- c(rep(0,2), rep(1,5), rep(0,3), rep(1,6),rep(0,2))
criteria1[c(6,9,12,13,14,15)] <- NA

#cumsum function, working before the first NA
ave(criteria1, cumsum(criteria1 == 0), FUN = cumsum )
[1]  0  0  1  1  1 NA  1  0 NA  0  1 NA NA NA NA  1  0  0

#and desired output would be
#NA's are replaced with the last value accumulated
#if more than three leave NA's in 
0 0 1 2 3 3 4 0 0 0 1 NA NA NA NA 2 0 0

Some conditions:

• NA s can not be replaced with zero(or one),
• vector must remain the same length (so excluding is not an option)
• longest length of consecutive ignored NA s should be three. If it is more than three, then they should remain as NA s and function should continue from the last non NA .

Some answers exist on the same topic, but I am not sure how to put it all together.
Thanks

Answer 1

With R base you can do: generate data

criteria1 <- c(rep(0,2), rep(1,5), rep(0,3), rep(1,6),rep(0,2))
criteria1[c(6,9,12,13)] <- NA

get result

l <- length(criteria1)
cum <- cumsum(ifelse(!is.na(criteria1),criteria1,0))
zero <- which(criteria1 == 0)

res <- cum - rep(cum[zero], c(zero[2:length(zero)],l+1)-zero)

optional dplyr solution:

res <- cum - rep(cum[zero], dplyr::coalesce(dplyr::lead(zero),l+1L)-zero)

detect and change repeats of NA > 3 times

NAs <- rle(is.na(criteria1))
NAloc <- which(NAs$lengths > 3 & NAs$values == 1)
for(i in NAloc)
{
res[seq(sum(NAs$lengths[1:(i-1)])+1,sum(NAs$lengths[1:i]))] <- NA
}

Answer 2

Since NA s are treated as zero when summed up but they are grouped as if they have the same value as previous values, you can treat NA differently based on the logic in the value variable and group variable within ave :

library(data.table); library(dplyr); library(zoo);

ave(coalesce(criteria1, 0), rleid(na.locf(criteria1 != 0)), FUN = cumsum)
# [1] 0 0 1 2 3 3 4 0 0 0 1 1 1 2 3 4 0 0