CODEWITHSUNDEEP
pythonarraysnumpy
camel_case
camel_case

Reputation: 23

How to progressively sum a NumPy array

Forgive me if this has been asked before, I couldn't find it. I am trying to progressively sum a numpy array into a new numpy array using vector operations. What I mean by this is that the 2nd index of the new array is equal to the 1st + 2nd index of the old array. or A[n] = B[0] + B[1] ... + B[n]. I know how to do this using a for loop but I'm looking for a vectorized solution.

Here is my non-vectorized solution: 

import numpy as np  
A = np.arange(10)
B = np.empty(10)

for i in range(len(A)):
    B[i] = sum(A[0:i+1])

print(B)

Upvotes: 2

Views: 363

Answers (2)

finefoot
finefoot

Reputation: 11302

The "progressive" sum is called cumulative sum. Use NumPy's cumsum for this.

Using your example and comparing B to np.cumsum(A) results in equal arrays:

>>> import numpy as np
>>> A = np.arange(10)
>>> B = np.empty(10)
>>> for i in range(len(A)):
...     B[i] = sum(A[0:i+1])
... 
>>> np.array_equal(B, np.cumsum(A))
True

Upvotes: 1

Wesley Yang
Wesley Yang

Reputation: 376

You can do it like this:

import numpy as np  

A = np.arange(10)
B = np.cumsum(A)
# [ 0  1  3  6 10 15 21 28 36 45]

Thanks

Upvotes: 5

Related Questions