Is there a NumPy function to compute partial sums?

Question

Given a 1-D array, A and a list of n unique and sorted indices idx, I would like to calculate sums of A[idx[i]:idx[i + 1]] for i = 0 to n - 1. A for-based solution is:

S = [A[idx[i]:idx[i + 1]].sum() for i in range(n - 1)]

But I suppose this would be very slow if n is large as it is done at Python level. Is there a NumPy function to achieve this (hopefully faster)?

Answer 1

You are looking for the reduceat method of np.add:

np.add.reduceat(A, idx)[:-1]

The [:-1] is just there to remove the last element that reduceat appends with the sum from idx[-1] to A.size.