Numpy: How to get the sums of integer array slices based on indexes from another array in a vectorized manner?

Question

Given two integer arrays a and b, where the elements in b represent indices of a...

a = array([10,10,10,8,8,8])

b = array([0,2,3,5])

I want to generate a new array whose elements are the sums of the elements in a along the index ranges given in b, not including the element at the tail of the range... It's hard to put in words, but the expected result given a and b from above would be:

result = array([0, # sum(a[:0])
                20,  # sum(a[0:2])
                10,  # sum(a[2:3])
                16])  # sum(a[3:5])

How could I achieve this in a vectorized/"numpythonic" way?

Thank you!

Answer 1

I think you are looking at np.ufunc.reduceat:

np.add.reduceat(a,b)

Out:

# gotta handle the case `b[0] == 0` separately
array([20, 10, 16,  8])

Answer 2

You can try this:

import numpy as np

a = np.array([10,10,10,8,8,8])
b = np.array([0,2,3,5])

list(map(sum, np.split(a, b)))

It gives:

[0, 20, 10, 16, 8]

The last number is the sum of the slice a[5:].