Itertools Recipe: Sliding Window

The itertools docs has a ton of slick recipes for using the library to good effect. Some of the code is more useful than illustrative, so I wanted to use these notebooks to break down a few of the functions.

# poor import style, but I want to copy-paste the code
# as-is from the docs

from itertools import *
import itertools

pairwise()

def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

This one’s super useful for generating sliding windows over an iterable

Demo

monotonically_increasing = lambda x: x[0] < x[1]

example = list(range(10000))

all(map(monotonically_increasing, pairwise(example)))

True

example = list(range(10000))
example[2001] = 0

all(map(monotonically_increasing, pairwise(example)))

False

Why this works

Simple enough, we get two of the exact same iterable, and then toss the first value of the second one away.

Then the zip() function will keep calling next() on each of them together until one of them runs out (the second one will, as it has one less value)

sliding_window()

If we wanted to extend the same functionality but across arbitrarily-many tee’d iterables, we can use the following

def sliding_window(iterable, n=2):
    iterables = itertools.tee(iterable, n)
    
    for iterable, num_skipped in zip(iterables, itertools.count()):
        for _ in range(num_skipped):
            next(iterable, None)
    
    return zip(*iterables)

Demo

example = range(10)
print(list(example))

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

for vals in sliding_window(example, 5):
    print(list(vals))

[0, 1, 2, 3, 4]
[1, 2, 3, 4, 5]
[2, 3, 4, 5, 6]
[3, 4, 5, 6, 7]
[4, 5, 6, 7, 8]
[5, 6, 7, 8, 9]

Why this works

Same general idea as above, the only difference here is using itertools.count() to dictate the number of leading values that we toss off of each tee’d iterator. Otherwise, this is a more-generic implementation than the above.