Nico Schlömer Nico Schlömer - 6 months ago 24
Python Question

NumPy: Select and sum data into array

I have a (large) data array and a (large) list of lists of (a few) indices, e.g.,

data = [1.0, 10.0, 100.0]
contribs = [[1, 2], [0], [0, 1]]

For each entry in
, I'd like to sum up the corresponding values of
and put those into an array. For the above example, the expected result would be

out = [110.0, 1.0, 11.0]

Doing this in a loop works,

c = numpy.zeros(len(contribs))
for k, indices in enumerate(contribs):
for idx in indices:
c[k] += data[idx]

but since
are large, it's taking too long.

I have the feeling this can be improved using numpy's fancy indexing.

Any hints?


One possibility would be

data = np.array(data)
out = [np.sum(data[c]) for c in contribs]

Should be faster than the double loop, at least.