cumsum in Julia with reset at 0

Question

Looking for an efficient solution that performs a cumulative sum that reset on zeros.

For example:

3×3 Matrix{Float64}:
 1.0  2.0  3.0
 0.0  5.0  6.0
10.0  0.0  9.0

the desired output is

3×3 Matrix{Float64}:
 1.0  2.0  3.0
 0.0  7.0  9.0
10.0  0.0  18.0

cumsum() exists but doesn't reset on zeros:

julia> cumsum(a, dims=1)
3×3 Matrix{Float64}:
  1.0  2.0   3.0
  1.0  7.0   9.0
 11.0  7.0  18.0

I have seen a lot of answers to this question in Python and R. I am looking for a solution in Julia, thank you.

Maybe you can comment with some answers from Python/R to compare? — Dan Getz, Jun 03 '23 at 22:19

score 4 · Accepted Answer · answered Jun 03 '23 at 22:14

julia> a = [
        1.0  2.0  3.0
        0.0  5.0  6.0
       10.0  0.0  9.0
       ]

julia> using BenchmarkTools

julia> @btime accumulate((a,b)->b == 0 ? b : a+b, $a; dims=1)
  50.466 ns (1 allocation: 128 bytes)
3×3 Matrix{Float64}:
  1.0  2.0   3.0
  0.0  7.0   9.0
 10.0  0.0  18.0

julia> @btime accumulate((a,b)->ifelse(b == 0, b, a+b), $a; dims=1)
  50.461 ns (1 allocation: 128 bytes)
3×3 Matrix{Float64}:
  1.0  2.0   3.0
  0.0  7.0   9.0
 10.0  0.0  18.0

julia> @btime cumsum($a.-($a .== 0).*cumsum($a, dims=1), dims=1)
  158.339 ns (3 allocations: 384 bytes)
3×3 Matrix{Float64}:
  1.0  2.0   3.0
  0.0  7.0   9.0
 10.0  0.0  18.0

The winning method is:

accumulate((a,b)->b == 0 ? b : a+b, a; dims=1)

A hair faster for big inputs (although it changes NaN/Inf behavior): `accumulate((s, x) -> muladd(!iszero(x), s, x), A; dims=1)` — August, Jun 03 '23 at 22:54

cumsum in Julia with reset at 0

1 Answers1