Cumulative sum while a condition is met with NA-CodePudding

Imagine I have these two vectors:

a <- c(0,0,1,1,1,1,1,2,2,2,2,2,2,2,3,3,3,3)
b <- c(NA,NA,NA,3,NA,NA,5,NA,NA,4,5,NA,2,NA,1,NA,NA,1)

And I am trying to have the cumulative sum by group that would end up in something like this:

c(NA,NA,NA,3,NA,NA,8,NA,NA,4,9,NA,11,NA,1,NA,NA,2)

I am trying with do.call(rbind,by(b,a,cumsum)) but it does not work, it returns an error

Warning message:
In (function (..., deparse.level = 1)  :
  number of columns of result is not a multiple of vector length (arg 1)

Any ideas? Thanks!

CodePudding user response：

You could use ave.

ave(b, a, FUN=\(x) {r <- cumsum(replace(x, is.na(x), 0)); replace(r, is.na(x), NA)})
# [1] NA NA NA  3 NA NA  8 NA NA  4  9 NA 11 NA  1 NA NA  2

CodePudding user response：

Another possible solution, based on dplyr and previous creation of a dataframe:

library(dplyr)

df <- data.frame(a, b)

df %>% 
  group_by(a) %>% 
  mutate(c = cumsum(ifelse(!is.na(b), b, 0))) %>% 
  mutate(c = ifelse(is.na(b), NA, c)) %>% 
  pull(c)

#>  [1] NA NA NA  3 NA NA  8 NA NA  4  9 NA 11 NA  1 NA NA  2