Search code examples
arraysmatlabmatrix-indexing

Vectorizing addition of subarray


Let's say I have two (large) vectors a=[0 0 0 0 0] and b=[1 2 3 4 5] of the same size and one index vector ind=[1 5 2 1] with values in {1,...,length(a)}. I would like to compute

for k = 1:length(ind)
    a(ind(k)) = a(ind(k)) + b(ind(k));
end
% a = [2 2 0 0 5]

That is, I want to add those entries of b declared in ind to a including multiplicity.

a(ind)=a(ind)+b(ind);
% a = [1 2 0 0 5]

is much faster, of course, but ignores indices which appear multiple times.

How can I speed up the above code?


Solution

  • We can use unique to identify the unique index values and use the third output to determine which elements of ind share the same index. We can then use accumarray to sum all the elements of b which share the same index. We then add these to the original value of a at these locations.

    [uniqueinds, ~, inds] = unique(ind);
    a(uniqueinds) = a(uniqueinds) + accumarray(inds, b(ind)).';
    

    If max(inds) == numel(a) then this could be simplified to the following since accumarray will simply return 0 for any missing entry in ind.

    a(:) = a(:) + accumarray(ind(:), b(ind));