List item arrangement in python

Viewed 110

I am looking for an efficient way to transform a given list into another list with degree n

here is input:

lst = [1, 2, 3, 4, 5, 6, 7]
n = 3

And favorable output is this:

[[1], [1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7], [7]]

It is basically like every element in resulting list is made of n consecutive elements from the given list except elements at starting and ending n-1 indices: 1, 2, 6, 7 in this example

Also for integer n, 1 <= n <= len(lst) is essential

1 Answers

You can use a comprehension with appropriate slices:

def chunks(lst, n=3):
    return [lst[max(i,0):i+n] for i in range(1-n, len(lst))]

chunks([1, 2, 3, 4, 5, 6, 7])
# [[1], [1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7], [7]]
chunks([1, 2, 3, 4, 5], 4)
# [[1], [1, 2], [1, 2, 3], [1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5], [4, 5], [5]]
chunks([1, 2], 2)
# [[1], [1, 2], [2]]
chunks([1, 2, 3], 1)
# [[1], [2], [3]]

A more elegant version of this sliding window uses a collections.deque (as pointed out by @chepner in the comments):

from collections import deque

def chunks(lst, n=3):
    window = deque(maxlen=n)
    for x in lst:
        window.append(x)
        yield list(window)
    while len(window) > 1:
        window.popleft()
        yield list(window)
  
Related