Search code examples
pythonnumpysumcumsum

Summing Two arrays : ary1.cumsum[:n] + ary2[n]?

Lets say I have two arrays of the same size, how can I in a numpy-way get single array with the following elements :

 a[:1].sum() + b[0] 
 a[:2].sum() + b[1]
 a[:3].sum() + b[2]
 a[:4].sum() + b[3]
 .....
Solution

  • a.cumsum() + b will do it.

    For example,

    In [11]: a
Out[11]: array([3, 5, 1, 4, 2, 6, 0, 5])

In [12]: b
Out[12]: array([1, 2, 3, 4, 5, 6, 7, 8])

In [13]: a.cumsum() + b
Out[13]: array([ 4, 10, 12, 17, 20, 27, 28, 34])

    Compare that to:

    In [14]: [a[:k+1].sum() + b[k] for k in range(len(a))]
Out[14]: [4, 10, 12, 17, 20, 27, 28, 34]