Vectorize summing values that belong to a group

2019-01-20 17:26发布

问题:

Is it possible to vectorize a loop that goes through different index mappings? For example:

a = zeros(1, 5);
m = [4 3 5; 5 1 3];
f = [1 2 3; 4 5 6];

for ii = 1:size(m,1)
    a(m(ii,:)) = a(m(ii,:)) + f(ii,:);
end

Gives output:

a = [5 0 2+6 1 3+4] = [5 0 8 1 7]

Can this be done without the for loop?

回答1:

This is a classic case of accumarray. accumarray works by providing a set of keys and a set of values associated with each key. accumarray groups all values that belong to the same key and does something to all of the values. The default behaviour is to sum all of the values that belong to the same key together, which is what you're after.

In your case, m are the keys and f are the values you want to add up that belong to the same key. Therefore:

>> a = accumarray(m(:), f(:))

a =

     5
     0
     8
     1
     7

In general, you may have keys that are missing. Therefore, you may opt to specify the output dimensions of the output array where it should be the maximum key value seen in m:

a = accumarray(m(:), f(:), [max(f(:)), 1]);

This is of course assuming that f consists of strictly positive values.

In general, if you have floating point numbers in f, then accumarray out of the box won't work because the keys are assumed to be strictly positive and integer. However, a common trick is to assign a unique ID to each value of f and use this as the input into accumarray. The third output of unique should do this for you. You'll also need the first output of unique to help you figure out which sum belongs to what key:

[msorted,~,id] = unique(m);
a = accumarray(id, f(:));
out = [msorted a];

out will contain a 2 column matrix where each row gives you a unique value in m and the associated sum for all values that shared the same key in m.