Numpy: How to get the sums of integer array slices based on indexes from another array in a vectoriz-CodePudding

Given two integer arrays a and b, where the elements in b represent indices of a...

a = array([10,10,10,8,8,8])

b = array([0,2,3,5])

I want to generate a new array whose elements are the sums of the elements in a along the index ranges given in b, not including the element at the tail of the range... It's hard to put in words, but the expected result given a and b from above would be:

result = array([0, # sum(a[:0])
                20,  # sum(a[0:2])
                10,  # sum(a[2:3])
                16])  # sum(a[3:5])

How could I achieve this in a vectorized/"numpythonic" way?

Thank you!

CodePudding user response：

I think you are looking at np.ufunc.reduceat:

np.add.reduceat(a,b)

Out:

# gotta handle the case `b[0] == 0` separately
array([20, 10, 16,  8])

CodePudding user response：

You can try this:

import numpy as np

a = np.array([10,10,10,8,8,8])
b = np.array([0,2,3,5])

list(map(sum, np.split(a, b)))

It gives:

[0, 20, 10, 16, 8]

The last number is the sum of the slice a[5:].

CodePudding user response：

Is that what you are looking for?

import numpy as np

a = np.array([10,10,10,8,8,8])

b = np.array([0,2,3,5])

result = []


for i, v in enumerate(b):
    result.append(sum(a[b[i-1]:v]))

result = np.array(result)

The result:

[ 0 20 10 16]