itertools: ready-made iteration tools

Ready-made iterator recipes

Every engineer eventually needs "all pairs of players", "these two lists glued together", or "the first 100 of an endless sequence". Writing those as nested index loops is slow to read and easy to get wrong, so the standard library ships them pre-built and pre-tested. itertools is that toolbox: lazy iterator builders, the same laziness you built by hand with generators in lesson 4-2.

These are the ones you will reach for most.

Because the results are lazy iterators, wrap them in list(...) to look at them, or loop over them directly.

islice(it, 5) stops asking after five items, which makes it the safe way to sample an infinite iterator like count. A plain list(count(10)) would try to build an endless list and never finish.

Building a round-robin schedule

combinations gives you every pair of players exactly once, never pairing anyone with themselves and never repeating a pairing in the other order, which is precisely what a round-robin schedule needs.

from itertools import combinations, chain, islice, count

players = ["mia", "leo", "zoe"]
for a, b in combinations(players, 2):
    print(a, "vs", b)

print(list(chain([1, 2], [3, 4])))
print(list(islice(count(100), 4)))

Output

mia vs leo
mia vs zoe
leo vs zoe
[1, 2, 3, 4]
[100, 101, 102, 103]

Three players produce three matches, not nine. The last line is the important one to notice: count(100) never ends, yet the program finishes, because islice stops requesting values once it has four.

Every size and color pairing

product produces every combination across two sequences, in the order nested loops would produce them, with the first argument as the outer loop.

from itertools import product

sizes = ["S", "M"]
colors = ["red", "blue"]

for size, color in product(sizes, colors):
    print(size, color)

Output

S red
S blue
M red
M blue

product(sizes, colors) is the same set of pairings you got from a two-clause comprehension in lesson 1-3, and it yields each result as a tuple, which the for line unpacks straight into size and color. Two sizes times two colors gives four rows, and the count multiplies out the same way for longer inputs.

Counting combinations against permutations

list(combinations([1, 2, 3, 4], 2)) contains 6 items: (1,2), (1,3), (1,4), (2,3), (2,4), and (3,4).

That is the binomial count C(4,2), the number of ways to choose 2 things from 4 when order does not matter. permutations([1, 2, 3, 4], 2) would give 12 instead, because there order does matter and (1,2) and (2,1) count separately.

Choosing teammates or pairs where the ordering carries no meaning: combinations. Choosing a first place and a second place: permutations.

Slicing an infinite stream

list(islice(count(1), 3)) evaluates to [1, 2, 3]. count(1) yields 1, 2, 3, and onward without end, and islice cuts the stream off after three items.

Laziness is the whole trick. The infinite iterator never computes a fourth value, because nothing ever asks for one. The same pattern lets you take a peek at the first few results of an expensive pipeline without paying to run all of it.