Numpy,将每个N连续元素分组?

时间:2015-04-26 09:03:57

标签: python arrays numpy

我想从数组中提取每N个连续元素的组。对于像这样的numpy数组:

a = numpy.array([1,2,3,4,5,6,7,8])

我希望(N = 5):

array([[1,2,3,4,5],
       [2,3,4,5,6],
       [3,4,5,6,7],
       [4,5,6,7,8]])

这样我就可以运行其他函数,例如average和sum。我如何制作这样的阵列?

2 个答案:

答案 0 :(得分:8)

使用broadcasting -

的一种方法
import numpy as np
out = a[np.arange(a.size - N + 1)[:,None] + np.arange(N)]

示例运行 -

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

In [32]: N
Out[32]: 5

In [33]: out
Out[33]: 
array([[4, 2, 5, 4, 1],
       [2, 5, 4, 1, 6],
       [5, 4, 1, 6, 7],
       [4, 1, 6, 7, 3]])

答案 1 :(得分:6)

您可以使用此blog

中的rolling_window
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]])

我喜欢@Divkar的解决方案。但是,对于较大的数组和窗口,您可能需要使用rolling_window

In [55]: a = np.arange(1000)

In [56]: %timeit rolling_window(a, 5)
100000 loops, best of 3: 9.02 µs per loop

In [57]: %timeit broadcast_f(a, 5)
10000 loops, best of 3: 87.7 µs per loop

In [58]: %timeit rolling_window(a, 100)
100000 loops, best of 3: 8.93 µs per loop

In [59]: %timeit broadcast_f(a, 100)
1000 loops, best of 3: 1.04 ms per loop