python - Sum array by number in numpy

Question

Assuming I have a numpy array like: [1,2,3,4,5,6] and another array: [0,0,1,2,2,1] I want to sum the items in the first array by group (the second array) and obtain n-groups results in group number order (in this case the result would be [3, 9, 9]). How do I do this in numpy?

Answer 1

The numpy function bincount was made exactly for this purpose and I'm sure it will be much faster than the other methods for all sizes of inputs:

data = [1,2,3,4,5,6]
ids  = [0,0,1,2,2,1]

np.bincount(ids, weights=data) #returns [3,9,9] as a float64 array

The i-th element of the output is the sum of all the data elements corresponding to "id" i.

Hope that helps.