Get sums of pairs of elements in a numpy array

Question

I have an array:

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

which is just a random shuffle of the range [0, 9]. I need to calculate this:

t2 = [9, 5, 7, 8, 7, 14, 11, 5, 11, 13]

which is just:

t2 = [t[0]+t[1], t[1]+t[2], t[2]+t[3], t[3]+t[4], ..., t[9]+t[0]]

Is there a way I can do this with numpy to avoid a python for loop when dealing with large arrays?

Answer 1

You could take advantage of a NumPy array's ability to sum element-wise:

In [5]: import numpy as np

In [6]: t = np.array([4, 5, 0, 7, 1, 6, 8, 3, 2, 9])

In [7]: t + np.r_[t[1:],t[0]]
Out[7]: array([ 9,  5,  7,  8,  7, 14, 11,  5, 11, 13])

np.r_ is one way to concatenate sequences together to form a new numpy array. As we'll see below, it turns out not to be the best way in this case.

---

Another possibility is:

In [10]: t + np.roll(t,-1)
Out[10]: array([ 9,  5,  7,  8,  7, 14, 11,  5, 11, 13])

It appears using np.roll is significantly faster:

In [11]: timeit t + np.roll(t,-1)
100000 loops, best of 3: 17.2 us per loop

In [12]: timeit t + np.r_[t[1:],t[0]]
10000 loops, best of 3: 35.5 us per loop

Answer 2

You can do this pretty happily with zip(), a list slice, and a list comprehension:

t2 = [a+b for (a, b) in zip(t, t[1:])]
t2.append(t[0]+t[-1])

We need the extra append() to add in the last element, as zip() only works until the shortest iterator ends. A list comprehension is significantly faster than a normal for loop as it's implemented C-side in Python, rather than as a Python loop.

The alternative is to use itertools.zip_longest:

from itertools import zip_longest
t2 = [a+b for (a, b) in zip_longest(t, t[1:], fillvalue=t[0])]

To fill the extra value in. Do note that this function is itertools.izip_longest in Python 2.x.

Answer 3

What about

import numpy as np
t = np.array([4, 5, 0, 7, 1, 6, 8, 3, 2, 9])

new_t = t + np.hstack((t[1:], [t[0]]))

Result:

>>> new_t
array([ 9,  5,  7,  8,  7, 14, 11,  5, 11, 13])