Question: Let us say you have a long list of integers, which you would like to iterate through N items at a time. How can you do that?
A concrete example is given below with N=3 (i.e. iterate three items at a time).
it = iter(range(1000))
# Write a program that will iterate through the it.
# For each iteration it should consume three elements:
# (0, 1, 2), (3, 4, 5), ...
First let us solve for the case where the number of items is exactly divided by N: len(it)%N == 0.
We can solve this case easily using zip() function. Zip takes a set of iterables as argument and produces an iterable that contains each element is a tuple with one element picked from the input list. For instance:
it1 = iter(range(5))
it2 = iter(range(5))
it3 = zip(it1, it2)
# it3 => (0, 0), (1, 1), (2, 2) ...
If we give the same iterator multiple times, zip will invoke the next(it) to form the tuple that it returns. Hence by passing the same argument twice, zip will be consuming two elements from the input iterable.
it1 = iter(range(10))
it3 = zip(it1, it1)
# it3 => (0, 1), (2, 3), ...
If we extend the idea, we need to produce N copies of the iterable as pass as argument to the zip function to consume N elements at a time.
it1 = iter(range(10)) #1
zip_arg = [it1 for _ in range(N)] #2
it3 = zip(*zip_arg) #3
# it3 => (0, 1, .. N-1), (N, N+1, .. 2N-1), ...
There are no surprises in #1. In #2, we are making N copies of the iterator using the list comprehension. In #3, we are de-referening the zip_arg so that it can be passed as argument to the zip function. Thats it.
What if the number of elements in the input list is not divisible by N? The last len(it)%N elements will skipped, as zip will stop with the shortest list. Example below shows what happens when N=2.
it1 = iter(range(5)) #1 (0, 1, 2, 3, 4)
zip_arg = [it1, it1] #2
it3 = zip(*zip_arg) #3
# it3 => (0, 1), (2, 3)
Thats where the itertools.zip_longest will get the job done. It will fill the result with None for shorter lists.
from itertools import zip_longest
it1 = iter(range(5)) #1 (0, 1, 2, 3, 4)
zip_arg = [it1, it1] #2
it3 = zip_longest(*zip_arg) #3
# it3 => (0, 1), (2, 3), (4, None)