dcrosta dcrosta - 2 months ago 5
Python Question

Iterate a list as pair (current, next) in Python

I sometimes need to iterate a list in Python looking at the "current" element and the "next" element. I have, till now, done so with code like:

for current, next in zip(the_list, the_list[1:]):
# Do something


This works and does what I expect, but is there's a more idiomatic or efficient way to do the same thing?

Answer

Here's a relevant example from the itertools module docs:

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

How this works:

First, two parallel iterators, a and b are created (the tee() call), both pointing to the first element of the original iterable. The second iterator, b is moved 1 step forward (the next(b, None)) call). At this point a points to s0 and b points to s1. Both a and b can traverse the original iterator independently - the izip function takes the two iterators and makes pairs of the returned elements, advancing both iterators at the same pace.

One caveat: the tee() function produces two iterators that can advance independently of each other, but it comes at a cost. If one of the iterators advances further than the other, then tee() needs to keep the consumed elements in memory until the second iterator comsumes them too (it cannot 'rewind' the original iterator). Here it doesn't matter because one iterator is only 1 step ahead of the other, but in general it's easy to use a lot of memory this way.