NumPy: Sum over 1-D array split by index

Question

Consider a 1-D NumPy input array and a sorted index array. The goal is to get the sum over the input array, a , but split by the indices defined in the index array.

Below are two approaches, but both of them require slow Python for loops. Is there a pure NumPy version without the need of Python for loops?

Example :

a = np.arange(20) # Input array
idxs = np.array([7, 15, 16]) # Index array

# Goal: Split a at index 7, 15 and 16 and
# compute sum for each partition

# Solution 1:
idxs_ext = np.concatenate(([0], idxs, [a.size]))
results = np.empty(idxs.size + 1)
for i in range(results.size):
    results[i] = a[idxs_ext[i]:idxs_ext[i+1]].sum()

# Solution 2:
result = np.array(
    [a_.sum() for a_ in np.split(a, idxs)]
)

# Result: array([21., 84., 15., 70.])

Answer 1

At first, you can split the a array based on your idxs array by np.split and then apply your function on that as:

np.stack(np.vectorize(np.sum)(np.array(np.split(a, idxs), dtype=object)))

Another answer is using np.add.reduceat as it is mentioned by @hpaulj in the comments and is more faster:

np.add.reduceat(a, np.insert(idxs, 0, 0), axis=0)