itertools ​

Python's built-in module itertools provides useful functions for working with iterables.

First, let's look at several "infinite" iterators provided by itertools:

import itertools
natuals = itertools.count(1)
for n in natuals:
    print(n)

This code will print the sequence of natural numbers endlessly, and you can only stop it by pressing Ctrl+C, because count() creates an infinite iterator.

The cycle() function infinitely repeats a given sequence:

import itertools
cs = itertools.cycle('ABC')  # Note that a string is a type of sequence
for c in cs:
    print(c)

This will also run indefinitely.

The repeat() function is used to repeat a single element indefinitely, but you can limit the number of repetitions by providing a second parameter:

ns = itertools.repeat('A', 3)
for n in ns:
    print(n)

This will print:

A
A
A

Infinite sequences only iterate indefinitely when used in a for loop; creating an iterator does not generate all infinite elements at once, as it's impossible to create infinite elements in memory.

While infinite sequences can iterate indefinitely, we often use functions like takewhile() to extract a finite sequence based on conditions:

natuals = itertools.count(1)
ns = itertools.takewhile(lambda x: x <= 10, natuals)
print(list(ns))

This will output:

[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

itertools also provides several useful iterator operation functions:

chain() ​

chain() can concatenate a series of iterators into a larger iterator:

for c in itertools.chain('ABC', 'XYZ'):
    print(c)

Output:

A
B
C
X
Y
Z

groupby() ​

groupby() groups adjacent duplicate elements in an iterator:

for key, group in itertools.groupby('AAABBBCCAAA'):
    print(key, list(group))

Output:

A ['A', 'A', 'A']
B ['B', 'B', 'B']
C ['C', 'C']
A ['A', 'A', 'A']

The grouping rule can be customized with a function. If two elements return the same value when passed to the function, they are considered part of the same group. To ignore case when grouping, you can use:

for key, group in itertools.groupby('AaaBBbcCAAa', lambda c: c.upper()):
    print(key, list(group))

Output:

A ['A', 'a', 'a']
B ['B', 'B', 'b']
C ['c', 'C']
A ['A', 'A', 'a']

Exercise ​

The value of π can be calculated using the formula:

[ \frac{4}{\pi} = 1 - \frac{1}{3} + \frac{1}{5} - \frac{1}{7} + \ldots ]

Using Python's itertools module, let's calculate the sum of the first N terms of this sequence:

import itertools

def pi(N):
    ' Calculate the value of pi '
    # Step 1: Create a sequence of odd numbers: 1, 3, 5, 7, 9, ...
    odds = itertools.count(1, 2)

    # Step 2: Take the first N terms: 1, 3, 5, ..., 2*N-1.
    terms = itertools.takewhile(lambda x: x < 2 * N, odds)

    # Step 3: Add positive and negative signs and divide by 4: 4/1, -4/3, 4/5, -4/7, ...
    series = ((4 if i % 2 == 0 else -4) / (1 + 2 * i) for i in range(N))

    # Step 4: Sum up:
    return sum(series)

# Tests:
print(pi(10))
print(pi(100))
print(pi(1000))
print(pi(10000))
assert 3.04 < pi(10) < 3.05
assert 3.13 < pi(100) < 3.14
assert 3.140 < pi(1000) < 3.141
assert 3.1414 < pi(10000) < 3.1415
print('ok')

Summary ​

The itertools module provides functions that handle iteration. Their return values are not lists but rather iterators, which only compute values when iterated over in a for loop.