Fast path for sliding_window() #871
Description
For small non-zero values of n, we can do better than the current implementation with:
def sliding_window(iterable, n):
"Collect data into overlapping fixed-length chunks or blocks."
# sliding_window('ABCDEFG', 4) → ABCD BCDE CDEF DEFG
iterators = tee(iterable, n)
for i, iterator in enumerate(iterators):
next(islice(iterator, i, i), None)
return zip(*iterators)
Performance:
- For small values of
n, the new code can be several times faster. - The time to yield the first output tuple is
O(n**2)wherenis the window size. - The remaining tuples are generated in
O(n)time.
New dispatch code:
def sliding_window(iterable, n):
if n > 0 and n <= 20:
return sliding_window_tee_islice_version(iterable, n)
return sliding_window_deque_version(iterable, n)
Timing code:
import collections
from itertools import tee, islice
def sliding_window1(iterable, n):
# Structure: no intermediate iterators
iterator = iter(iterable)
window = collections.deque(islice(iterator, n - 1), maxlen=n)
for x in iterator:
window.append(x)
yield tuple(window)
def sliding_window2(iterable, n):
# Structure: zip of tees
iterators = tee(iterable, n)
for i, iterator in enumerate(iterators):
next(islice(iterator, i, i), None)
return zip(*iterators)
def sliding_window3(iterable, n):
# Structure: zip of islices of tees
# Always worse than sliding_window2
return zip(*(islice(t, i, None) for i, t in enumerate(tee(iterable, n))))
def sliding_window4(iterable, n):
# Structure: single tuple with no intermediate iterators
iterator = iter(iterable)
t = tuple(islice(iterator, n))
if len(t) == n:
yield t
for x in iterator:
yield (t := t[1:] + (x,))
# Smoke test
seq, k, target = 'ABCDEFG', 4, ['ABCD', 'BCDE', 'CDEF', 'DEFG']
assert list(map(''.join, sliding_window1(seq, k))) == target
assert list(map(''.join, sliding_window2(seq, k))) == target
assert list(map(''.join, sliding_window3(seq, k))) == target
assert list(map(''.join, sliding_window4(seq, k))) == target
data = [None] * 10_000
if __name__ == '__main__':
from timeit import repeat
setup = 'from __main__ import data, sliding_window1, sliding_window2, sliding_window3, sliding_window4'
stmts = [
'for w in sliding_window1(data, 4): pass',
'for w in sliding_window2(data, 4): pass',
'for w in sliding_window3(data, 4): pass',
'for w in sliding_window4(data, 4): pass',
'pass',
'for w in sliding_window1(data, 400): pass',
'for w in sliding_window2(data, 400): pass',
'for w in sliding_window3(data, 400): pass',
'for w in sliding_window4(data, 400): pass',
]
n = 10 ** 2
for stmt in stmts:
timing = min(repeat(stmt, setup, number=n, repeat=7))
print(f'{timing:7.4f} {stmt!r}')
Timings with Python 3.13b3 running on an Apple M1 Max:
0.0721 'for w in sliding_window1(data, 4): pass'
0.0237 'for w in sliding_window2(data, 4): pass'
0.0334 'for w in sliding_window3(data, 4): pass'
0.0656 'for w in sliding_window4(data, 4): pass'
0.0000 'pass'
1.2481 'for w in sliding_window1(data, 400): pass'
1.2173 'for w in sliding_window2(data, 400): pass'
2.1797 'for w in sliding_window3(data, 400): pass'
1.1178 'for w in sliding_window4(data, 400): pass'