Numpy: How to get the sums of integer array slices based on indexes from another array in a vectorized manner?

Question

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!

Quang Hoang · Accepted Answer · 2021-10-28 23:14:20Z

5

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

answered Oct 28, 2021 at 23:14

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

bb1 · Accepted Answer · 2021-10-28 23:04:39Z

2

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

answered Oct 28, 2021 at 23:04

bb1

7,9332 gold badges11 silver badges26 bronze badges

1 Comment

Quang Hoang Over a year ago

This, while being very concise, is not vectorized.

Pedro Maia · Accepted Answer · 2021-10-28 23:09:15Z

0

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]

edited Oct 28, 2021 at 23:09

answered Oct 28, 2021 at 23:00

Pedro Maia

2,7321 gold badge7 silver badges21 bronze badges

1 Comment

Alex Parthemer Over a year ago

Almost, Except I'm wanting to avoid iteration for performance reasons. In practice, my arrays are much larger. Also, the first element of the result array should be zero per the expected result given in the question. Thank you though!

Collectives™ on Stack Overflow

Numpy: How to get the sums of integer array slices based on indexes from another array in a vectorized manner?

3 Answers 3

Comments

1 Comment

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

1 Comment

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related