CODEWITHSUNDEEP
pythonnumpysplitsumvectorization
Yannic
Yannic

Reputation: 817

NumPy: Sum over 1-D array split by index

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.])

Upvotes: 2

Views: 259

Answers (1)

Ali_Sh
Ali_Sh

Reputation: 2816

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)

Upvotes: 3

Related Questions