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.])
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)
Update:
Using np.concatenate instead of insert reduced the runtime 5 times for data range 1000 with 7 slices; It was the fastest method I have examined:
np.add.reduceat(a, np.concatenate(([0], idxs)), axis=0)