sum numpy array at given indices

Question

I want to add the values of a vector:

a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype='d')

to the values of another vector:

c = np.array([10, 10, 10], dtype='d')

at position given by another array (of the same size of a, with values 0 <= b[i] < len(c))

b = np.array([2, 0, 1, 0, 2, 0, 1, 1, 0, 2], dtype='int32')

This is very simple to write in pseudo code:

for I in range(b.shape[0]):
    J = b[I]
    c[J] += a[I]

Something like this, but vectorized (length of c is some hundreds in real case).

c[0] += np.sum(a[b==0]) # 27 (10 + 1 + 3 + 5 + 8)
c[1] += np.sum(a[b==1]) # 25 (10 + 2 + 6 + 7)
c[2] += np.sum(a[b==2]) # 23 (10 + 0 + 4 + 9)

My initial guess was:

c[b] += a

but only last values of a are summed.

Answer 1

You can use np.bincount to get ID based weighted summations and then add with c, like so -

np.bincount(b,a) + c