CODEWITHSUNDEEP
pythonarraysnumpy
Bill
Bill

Reputation: 700

Sum values according to an index array

I have two arrays of the same dimension:

a = np.array([ 1,  1, 2,  0, 0, 1])
b = np.array([50, 51, 6, 10, 3, 2])

I want to sum the elements of b according to the indices in a.

The ith element of the matrix I want will be the sum of all values b[j] such that a[j]==i.

So the result should be a 3-dimensional array of [10 + 3, 50 + 51 + 2, 6]

Is there a numpy way to do this? I have some very large arrays that I need to sum like this over multiple dimensions, so it would NOT be convenient to to have to perform explicit loops.

Upvotes: 2

Views: 2341

Answers (2)

njzk2
njzk2

Reputation: 39386

In case you are not using numpy, something as simple as :

res = [0]*len(set(a))
for i, v in enumerate(b):
    res[a[i]] += v

Assuming the indices in a are always 0-based and a continuous sequence.

Upvotes: 2

unutbu
unutbu

Reputation: 879103

numpy.bincount has a weights parameter which does just what you need:

In [36]: np.bincount(a, weights=b)
Out[36]: array([  13.,  103.,    6.])

Upvotes: 10

Related Questions