Rob Rob - 1 month ago 8
Python Question

Sum only the previous two values in numpy

If I have a list like [1, 2, 3, 4, 1, 2, 3, 4] and I wanted to sum the values like:

[6, 9, 8, 6, 7, 4]


where the formula is
value + value_one_ahead + value_two_ahead
.

Answer

The operation can be interpreted as a (discrete) convolution of the input with [1, 1, 1], so you can use numpy.convolve:

In [178]: x = np.array([1, 2, 3, 4, 1, 2, 3, 4])

In [179]: np.convolve(x, np.ones(3, dtype=x.dtype), mode='valid')
Out[179]: array([6, 9, 8, 7, 6, 9])

See the docstring for the description of the mode argument. You might have to tweak the return value or add a line or two of additional code if the available mode options don't handle the ends of the input the way you want. For example, to get something closer to what you wrote in the question, you can use mode='full' (which is the default), and then discard the first two elements (by taking the slice [2:]):

In [181]: np.convolve(x, np.ones(3, dtype=x.dtype), mode='full')[2:]
Out[181]: array([6, 9, 8, 7, 6, 9, 7, 4])

Note: np was defined using import numpy as np. np.ones(3, dtype=x.dtype) creates an array of 1s with the given length and data type:

In [180]: np.ones(3, dtype=x.dtype)
Out[180]: array([1, 1, 1])