Numpy, grouping every N continuous element?
One approach with broadcasting
-
import numpy as npout = a[np.arange(a.size - N + 1)[:,None] + np.arange(N)]
Sample run -
In [31]: aOut[31]: array([4, 2, 5, 4, 1, 6, 7, 3])In [32]: NOut[32]: 5In [33]: outOut[33]: array([[4, 2, 5, 4, 1], [2, 5, 4, 1, 6], [5, 4, 1, 6, 7], [4, 1, 6, 7, 3]])
You could use rolling_window
from this blog
def rolling_window(a, window): shape = a.shape[:-1] + (a.shape[-1] - window + 1, window) strides = a.strides + (a.strides[-1],) return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)In [37]: a = np.array([1,2,3,4,5,6,7,8])In [38]: rolling_window(a, 5)Out[38]:array([[1, 2, 3, 4, 5], [2, 3, 4, 5, 6], [3, 4, 5, 6, 7], [4, 5, 6, 7, 8]])
I liked @Divkar's solution. However, for larger arrays and windows, you may want to use rolling_window
?
In [55]: a = np.arange(1000)In [56]: %timeit rolling_window(a, 5)100000 loops, best of 3: 9.02 µs per loopIn [57]: %timeit broadcast_f(a, 5)10000 loops, best of 3: 87.7 µs per loopIn [58]: %timeit rolling_window(a, 100)100000 loops, best of 3: 8.93 µs per loopIn [59]: %timeit broadcast_f(a, 100)1000 loops, best of 3: 1.04 ms per loop