Sums of variable size chunks of a list where sizes are given by other list

user_13491274 0
Viewed 223

I would like to make the following sum given two lists:

a = [0,1,2,3,4,5,6,7,8,9]
b = [2,3,5]

The result should be the sum of the every b element of a like:

  1. b[0] = 2 so the first sum result should be: sum(a[0:2])
  2. b[1] = 3 so the second sum result should be: sum(a[2:5])
  3. b[2] = 5 so the third sum result should be: sum(a[5:10])

The printed result: 1,9,35

numpy python
5 Answers

You can make use of np.bincount with weights:

groups = np.repeat(np.arange(len(b)), b)

np.bincount(groups, weights=a)

Output:

array([ 1.,  9., 35.])
user_42384080

NumPy has a tool to do slice based sum-reduction with np.add.reduceat -

In [46]: np.add.reduceat(a,np.cumsum(np.r_[0,b[:-1]]))                          
Out[46]: array([ 1,  9, 35])
user_32938810

You mean something like this?

a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
b = [2, 3, 5]

def some_function(a, b): # couldnt come up with a name :D
    last_index = 0

    for i in b:
        print(sum(a[last_index:last_index + i]))
        last_index += i

some_function(a, b)
user_113603180

Hard to compete with the np.bincount solution, but here's another nice way to approach it with np.cumsum:

strides = [0] + np.cumsum(b).tolist()  # [0, 2, 5, 10]
stride_slices = zip(strides[:-1], strides[1:])  # [(0, 2), (2, 5), (5, 10)]
[sum(a[s[0]: s[1]]) for s in stride_slices]
# [1, 9, 35]
user_30157340

You can use a list comprehension with sum:

a=[0,1,2,3,4,5,6,7,8,9]
b=[2,3,5]
r = [sum(a[(k:=sum(b[:i])):k+j]) for i, j in enumerate(b)]

Output:

[1, 9, 35]
user_73267380
Related