# How to compute a conditional sum of a vector efficiently

**URL:** <https://discourse.julialang.org/t/how-to-compute-a-conditional-sum-of-a-vector-efficiently/85699>\
**Category:** General Usage\
**Created:** [August 13, 2022, 9:08am UTC](https://discourse.julialang.org/t/how-to-compute-a-conditional-sum-of-a-vector-efficiently/85699 "2022-08-13T09:08:55Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![structural](https://avatars.discourse-cdn.com/v4/letter/s/34f0e0/32.png) [@structural](https://discourse.julialang.org/u/structural)\
**Post date:** [August 13, 2022, 9:08am UTC](https://discourse.julialang.org/t/how-to-compute-a-conditional-sum-of-a-vector-efficiently/85699/1 "2022-08-13T09:08:55Z")

</div>

I have three vectors `x, a, b`. I need a new vector `Y` which has as its k^{th} element the following sum: Y^{}\_k = \sum\_{i=1}^N x\_i \ \mathbb{1}(a\_i=a\_k, b\_i=b\_k) i.e., the sum of `x` from all rows with the same values of `a, b` as the current row.

Is there a more efficient way to implement this than the code below? Would a for loop  
be better?

```julia
condition = (a.==a').*(b.==b')
y = (condition.*x', dims=2)

```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/224680_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [August 13, 2022, 10:44am UTC](https://discourse.julialang.org/t/how-to-compute-a-conditional-sum-of-a-vector-efficiently/85699/2 "2022-08-13T10:44:59Z")

</div>

I thought your code would be easy to optimize because the first line allocates a matrix of size n^2. My loop implementation is indeed faster in some cases, but slower in others, like the one below. Wonder what I’m doing wrong…

```julia
julia> using BenchmarkTools

julia> N, V = 1000, 3;

julia> a, b = rand(1:V, N), rand(1:V, N);

julia> x = rand(N);

julia> function f(x, a, b)
           condition = (a .== a') .* (b .== b')
           y = sum(condition .* x', dims=2)
       end
f (generic function with 1 method)

julia> function g(x, a, b)
           y = zero(x)
           for k in eachindex(x, a, b)
               aₖ, bₖ = a[k], b[k]
               for i in eachindex(x, a, b)
                   if (a[i] == aₖ) && (b[i] == bₖ)
                       y[k] += x[i]
                   end
               end
           end
           return y
       end
g (generic function with 1 method)

julia> @assert f(x, a, b) ≈ g(x, a, b)

julia> @btime f($x, $a, $b);
  2.371 ms (11 allocations: 7.76 MiB)

julia> @btime g($x, $a, $b);
  3.387 ms (1 allocation: 7.94 KiB)

```

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [August 13, 2022, 11:53am UTC](https://discourse.julialang.org/t/how-to-compute-a-conditional-sum-of-a-vector-efficiently/85699/3 "2022-08-13T11:53:35Z")

</div>

> [@gdalle](#):
>
> ```julia-auto
> julia> N, V = 1000, 3;
> 
> julia> a, b = rand(1:V, N), rand(1:V, N);
> 
> ```

If you only have a small number `V` of possible values of `a` and `b`, then you can do _much_ better by first summing into a matrix indexed by values of `V`, because then the time complexity is \Theta(N) (with space complexity \Theta(N+V^2)) instead of \Theta(N^2) as in the implementations above. Correspondingly, the following is about 1000x faster than either of the implementations above on my machine for this test data:

```julia-auto
function h(x, a, b, V)
    s = zeros(eltype(x), V, V)
    for k in eachindex(x, a, b)
        s[a[k], b[k]] += x[k]
    end
    return [s[a[k], b[k]] for k in eachindex(x, a, b)]
end

```

If you don’t know ahead of time what possible values there are for `a` and `b` (or if the number of possible values is large), you could use the same strategy but with a dictionary of tuples:

```julia-auto
function h2(x, a, b)
    s = Dict{Tuple{eltype(a),eltype(b)},eltype(x)}()
    for k in eachindex(x, a, b)
        key = (a[k], b[k])
        s[key] = get(s, key, zero(eltype(x))) + x[k]
    end
    return [s[a[k], b[k]] for k in eachindex(x, a, b)]
end

```

which also has \Theta(N) complexity, with a worse constant factor by about 6x on my machine. (Could be made a bit faster ala [julia#12157](https://github.com/JuliaLang/julia/issues/12157).)

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [August 13, 2022, 6:56pm UTC](https://discourse.julialang.org/t/how-to-compute-a-conditional-sum-of-a-vector-efficiently/85699/4 "2022-08-13T18:56:18Z")

</div>

2 posts were split to a new topic: [Speeding up a sum involving 3 matrices](https://discourse.julialang.org/t/speeding-up-a-sum-involving-3-matrices/85713)