Cumulative summation of a numpy array by index

Question

Assume you have an array of values that will need to be summed together

d = [1,1,1,1,1]

and a second array specifying which elements need to be summed together

i = [0,0,1,2,2]

The result will be stored in a new array of size max(i)+1. So for example i=[0,0,0,0,0] would be equivalent to summing all the elements of d and storing the result at position 0 of a new array of size 1.

I tried to implement this using

c = zeros(max(i)+1)
c[i] += d

However, the += operation adds each element only once, thus giving the unexpected result of

[1,1,1]

instead of

[2,1,2]

How would one correctly implement this kind of summation?

Answer 1

If I understand the question correctly, there is a fast function for this (as long as the data array is 1d)

>>> i = np.array([0,0,1,2,2])
>>> d = np.array([0,1,2,3,4])
>>> np.bincount(i, weights=d)
array([ 1.,  2.,  7.])

np.bincount returns an array for all integers range(max(i)), even if some counts are zero

Answer 2

Juh_'s comment is the most efficient solution. Here's working code:

import numpy as np
import scipy.ndimage as ni

i = np.array([0,0,1,2,2])
d = np.array([0,1,2,3,4])

n_indices = i.max() + 1
print ni.sum(d, i, np.arange(n_indices))

Answer 3

This solution should be more efficient for large arrays (it iterates over the possible index values instead of the individual entries of i):

import numpy as np

i = np.array([0,0,1,2,2])
d = np.array([0,1,2,3,4])

i_max = i.max()
c = np.empty(i_max+1)
for j in range(i_max+1):
    c[j] = d[i==j].sum()

print c
[1. 2. 7.]

Answer 4

def zeros(ilen):
 r = []
 for i in range(0,ilen):
     r.append(0)

i_list = [0,0,1,2,2]
d = [1,1,1,1,1]
result = zeros(max(i_list)+1)

for index in i_list:
  result[index]+=d[index]

print result