Python sliding windows of a list

Question

Is there an efficient or elegant way to retrieve all the k-size sublists of a list in Python? For example:

arr = [2, 3, 5, 7, 11, 13]

I want all 3-element sublists:

result = [[2, 3, 5],
          [3, 5, 7],
          [5, 7, 11],
          [7, 11, 13]]

I know I could create this with a for loop, slicing the list with arr[i:i+3], but the lists I'm dealing with are gigantic and I'm hoping for an efficient mechanism, or at least an elegant or Pythonic mechanism.

I'm using Pandas as well, so happy to use a Pandas mechanism.

Answer 1

If you actually want to construct the list, I don't think you'll do better than a basic list comprehension like this:

arr = [2, 3, 5, 7, 11, 13]
result = [arr[i:i+k] for i in range(len(arr)-k+1)]

If you want to minimize memory use, you could use a generator:

arr = [2, 3, 5, 7, 11, 13]
def window(arr, k):
    for i in range(len(arr)-k+1):
        yield arr[i:i+k]

for group in window(arr, 3):
    ...  # do something with group

You could also do something where you zip together k copies of the list, each offset by one. But that would take as much memory as the first solution, probably without much performance advantage.

There may be something quick and efficient in numpy or pandas, but you would need to show more about what your input and output should look like.

There are some other ideas here, but they are focused on general iterables (where you can only pull items out once), rather than lists (where you can access items by index, possibly repeatedly).

Answer 2

You can use strides:

arr = [2, 3, 5, 7, 11, 13]

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)

a = rolling_window(np.array(arr), 3)
print (a)
[[ 2  3  5]
 [ 3  5  7]
 [ 5  7 11]
 [ 7 11 13]]


print (a.tolist())
[[2, 3, 5], 
 [3, 5, 7], 
 [5, 7, 11], 
 [7, 11, 13]]