I need a rolling window (aka sliding window) iterable over a sequence/iterator/generator. (Default Python iteration could be considered a special case, where the window length is 1.) I'm currently using the following code. How can I do it more elegantly and/or efficiently?
def rolling_window(seq, window_size):
it = iter(seq)
win = [it.next() for cnt in xrange(window_size)] # First window
yield win
for e in it: # Subsequent windows
win[:-1] = win[1:]
win[-1] = e
yield win
if __name__=="__main__":
for w in rolling_window(xrange(6), 3):
print w
"""Example output:
[0, 1, 2]
[1, 2, 3]
[2, 3, 4]
[3, 4, 5]
"""
For the specific case of window_size == 2 (i.e., iterating over adjacent, overlapping pairs in a sequence), see also How can I iterate over overlapping (current, next) pairs of values from a list?.
There's one in an old version of the Python docs with itertools examples:
from itertools import islice
def window(seq, n=2):
"Returns a sliding window (of width n) over data from the iterable"
" s -> (s0,s1,...s[n-1]), (s1,s2,...,sn), ... "
it = iter(seq)
result = tuple(islice(it, n))
if len(result) == n:
yield result
for elem in it:
result = result[1:] + (elem,)
yield result
The one from the docs is a little more succinct and uses itertools to greater effect I imagine.
If your iterator is a simple list/tuple a simple way to slide through it with a specified window size would be:
seq = [0, 1, 2, 3, 4, 5]
window_size = 3
for i in range(len(seq) - window_size + 1):
print(seq[i: i + window_size])
Output:
[0, 1, 2]
[1, 2, 3]
[2, 3, 4]
[3, 4, 5]