Is there a NumPy function to compute partial sums?

user_1208923 0
Viewed 141

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)?

python-3.x numpy python
1 Answers

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.

user_29887300
Related