A few hours back I stumbled into a problem where I had to perform a lookahead of n elements in a list to do some calculations. The first thought: Just take the current index and get all elements until i+n. I started writing..
for i in range(len(iterable)): ---- SNAP ----
Stop. This is awfully unpythonic. There has to be a better way!
Browsing the itertools recipes for inspiration, I found the pairwise
function:
def pairwise(iterable): "s -> (s0,s1), (s1,s2), (s2, s3), ..." a, b = tee(iterable) next(b, None) return zip(a, b)
Perfect! Now I just have to adjust it to work with n iterators where the first iterator is at i, the second at i+1, etc.
def nwise(iterable, n=1): iterators = tee(iterable, n) for i in range(n): for _ in range(i): next(iterators[i], None) return zip(*iterators)
There we go. Even though it works on multiple levels of the iterable, it’s still memory-efficient, because generators are awesome. This will provide you with an output like this:
In [4]: l = [1,2,3,4,5,6] In [5]: list(nwise(l, n=3)) Out[5]: [(1, 2, 3), (2, 3, 4), (3, 4, 5), (4, 5, 6)]
Note that the list call was just used to empty the generator for printing. Here’s a quick one-liner that counts the times a fixed-length (42) sequence in a list sums up to a certain value (1337):
sum([1 for seq in nwise(l, n=42) if sum(seq) == 1337])