Search code examples
pythonarraysnumpysumsliding-window

Sum elements in window by n step in Numpy


Is a Numpy way to make a sum each three elements in the interval with a step? For example:

import numpy as np
mydata = np.array([4, 2, 3, 8, -6, 10])

I would like to get this result:

np.array([9, 12])

I suppose that np.convole can do this, according to Summing elements in a sliding window - NumPy, but can I change the step from n=1 to n=3, in this case?


Solution

  • Using reshape (requires the array length to be a multiple of n):

    n = 3
    
    mydata.reshape(-1, n).sum(1)
    

    If you don't have a multiple, you can trim:

    n = 3
    
    mydata[:len(mydata)//n*n].reshape(-1, n).sum(1)
    

    Using convolve, which should be much less efficient for large n as many values (n-1 out of n) are computed for nothing:

    np.convolve(np.ones(n), mydata)[n-1::n]
    

    Output: array([ 9, 12])