Cumulative sum that resets when 0 is encountered

Question

I would like to do a cumulative sum on a field but reset the aggregated value whenever a 0 is encountered.

Here is an example of what I want :

data.frame(campaign = letters[1:4] , 
       date=c("jan","feb","march","april"),
       b = c(1,0,1,1) ,
       whatiwant = c(1,0,1,2)
       )

 campaign  date b whatiwant
1        a   jan 1         1
2        b   feb 0         0
3        c march 1         1
4        d april 1         2

Answer 1

Another base would be just

with(df, ave(b, cumsum(b == 0), FUN = cumsum))
## [1] 1 0 1 2

This will just divide column b to groups according to 0 appearances and compute the cumulative sum of b per these groups

---

Another solution using the latest data.table version (v 1.9.6+)

library(data.table) ## v 1.9.6+
setDT(df)[, whatiwant := cumsum(b), by = rleid(b == 0L)]
#    campaign  date b whatiwant
# 1:        a   jan 1         1
# 2:        b   feb 0         0
# 3:        c march 1         1
# 4:        d april 1         2

---

Some benchmarks per comments

set.seed(123)
x <- sample(0:1e3, 1e7, replace = TRUE)
system.time(res1 <- ave(x, cumsum(x == 0), FUN = cumsum))
# user  system elapsed 
# 1.54    0.24    1.81 
system.time(res2 <- Reduce(function(x, y) if (y == 0) 0 else x+y, x, accumulate=TRUE))
# user  system elapsed 
# 33.94    0.39   34.85 
library(data.table)
system.time(res3 <- data.table(x)[, whatiwant := cumsum(x), by = rleid(x == 0L)])
# user  system elapsed 
# 0.20    0.00    0.21 

identical(res1, as.integer(res2))
## [1] TRUE
identical(res1, res3$whatiwant)
## [1] TRUE

Answer 2

Another late idea:

ff = function(x)
{
    cs = cumsum(x)
    cs - cummax((x == 0) * cs)
}
ff(c(0, 1, 3, 0, 0, 5, 2))
#[1] 0 1 4 0 0 5 7

And to compare:

library(data.table)
ffdt = function(x) 
    data.table(x)[, whatiwant := cumsum(x), by = rleid(x == 0L)]$whatiwant

x = as.numeric(x) ##because 'cumsum' causes integer overflow
identical(ff(x), ffdt(x))
#[1] TRUE
microbenchmark::microbenchmark(ff(x), ffdt(x), times = 25)
#Unit: milliseconds
#    expr      min       lq   median       uq      max neval
#   ff(x) 315.8010 362.1089 372.1273 386.3892 405.5218    25
# ffdt(x) 374.6315 407.2754 417.6675 447.8305 534.8153    25

Answer 3

You could use the Reduce function with a custom function that returns 0 when the new value encountered is 0 and otherwise adds the new value to the accumulated value:

Reduce(function(x, y) if (y == 0) 0 else x+y, c(1, 0, 1, 1), accumulate=TRUE)
# [1] 1 0 1 2

Answer 4

hutilscpp::cumsum_reset is designed for this purpose. The first argument is a logical vector, indicating when the cumulative sum should continue. The second argument is the input to the cumulative sum itself

library(hutilscpp)
b <- c(1, 0, 1, 1)
cumsum_reset(as.logical(b), b)

On my machine, compared to the data.table function above, this use of cumsum_reset is about 3 times faster.

Answer 5

Another variant using rep of the values at places where to reset. This uses cumsum only once but the drawback will lead to large numbers and can cause an integer overflow or inaccurate numeric values and will propagate NA to all following groups.

x <- cumsum(DF$b)
i <- which(DF$b == 0)
x - rep(c(0, x[i]), diff(c(1L, i, length(x)+1L)))
#[1] 1 0 1 2

Another way is to use Rcpp - in this case for integer.

Rcpp::cppFunction('IntegerVector csrA(const IntegerVector x, int z=0) {
  IntegerVector out(no_init(x.size()));
  int init = z == NA_INTEGER ? 0 : z;
  int s = 0;
  for(int i = 0; i < x.size(); ++i) {
    if(x[i] == z) s = init;
    else [[likely]] s += x[i];
    out[i] = s;
  }
  return out;
}')
csrA(DF$b)
#[1] 1 0 1 2

A variant also taking care of NA might look like:

Rcpp::cppFunction('IntegerVector csr(const IntegerVector x, int z=0) {
  IntegerVector out(no_init(x.size()));
  int init = z == NA_INTEGER ? 0 : z;
  LogicalVector isNA = is_na(x);
  int s = 0;
  for(int i = 0; i < x.size(); ++i) {
    if(x[i] == z) s = init;
    else [[likely]] if(isNA[i] || s == NA_INTEGER) s = NA_INTEGER;
      else [[likely]] s += x[i];
    out[i] = s;
  }
  return out;
}')

csr(c(2,4,3,0,3,5), 0)
#[1] 2 6 9 0 3 8

csr(c(2,NA,3,0,3,5), 0)
#[1]  2 NA NA  0  3  8

csr(c(2,4,3,1,3,5), 1)
#[1] 2 6 9 1 4 9

csr(c(2,4,3,NA,3,5), NA)
#[1] 2 6 9 0 3 8

Data

DF <- data.frame(campaign = letters[1:4] , 
                 date=c("jan","feb","march","april"),
                 b = c(1,0,1,1) ,
                 whatiwant = c(1,0,1,2)
                 )

---

Benchmark - Based on @David Arenburg

set.seed(123)
#Using 1e3 instead of 1e2 would lead to an integer overflow for whichRep and cummax
x <- sample(0:1e2, 1e7, TRUE)

library(data.table)

bench::mark(
ave = ave(x, cumsum(x == 0), FUN = cumsum),
data.table = data.table(x)[, whatiwant := cumsum(x), by = rleid(x == 0L)]$whatiwant,
cummax = {cs = cumsum(x)
  cs - cummax((x == 0) * cs)},
whichRep = {y <- cumsum(x)
i <- which(x == 0)
y - rep(c(0, y[i]), diff(c(1L, i, length(x)+1L)))},
RcppNA = csr(x),
RcppSimple = csrA(x)
)

Result

  expression      min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_…¹
  <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl> <bch:tm>
1 ave           1.06s    1.06s     0.945   751.8MB     3.78     1     4    1.06s
2 data.table 199.01ms 266.26ms     3.76    231.9MB     1.88     2     1 532.53ms
3 cummax      90.57ms  93.76ms    10.4     152.6MB     6.92     6     4  578.4ms
4 whichRep     74.5ms  77.05ms    12.9     195.6MB    11.1      7     6 541.63ms
5 RcppNA      39.55ms  40.84ms    24.2      76.3MB     5.60    13     3  536.1ms
6 RcppSimple  29.73ms  30.59ms    32.3      38.1MB     3.80    17     2  526.1ms