Nico Schlömer Nico Schlömer - 5 months ago 19
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
contribs
, I'd like to sum up the corresponding values of
data
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
data
and
contribs
are large, it's taking too long.

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

Any hints?

Answer

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.

Comments