How to speed up rowsum function?

General Usage Performance
3

Here’s one way to improve this, but there might be better ways:

julia> function rowsum2(mat::AbstractMatrix{<:Number}, groups::AbstractVector{<:Integer})
           size(mat, 1) == length(groups) || error("length of group vector must match matrix")
           rst = fill!(similar(mat, maximum(groups), size(mat, 2)), 0)
           for (i, g) in pairs(groups)
             @views rst[g, :] .+= mat[i, :]
           end
           return rst
       end;

julia> r = @btime rowsum($a, $g);
  min 1.125 s, mean 1.127 s (160002 allocations, 331.12 MiB)

julia> r ≈ @btime rowsum2($a, $g)
  min 984.500 μs, mean 1.061 ms (2 allocations, 625.05 KiB)
true
2 Likes