Python’s built-in module itertools provides extremely useful functions for manipulating iterable objects.
First, let’s look at several “infinite” iterators provided by itertools:
>>> import itertools
>>> natuals = itertools.count(1)
>>> for n in natuals:
... print(n)
...
1
2
3
...
Since count() creates an infinite iterator, the above code will print the sequence of natural numbers and never stop—you can only exit by pressing Ctrl+C.
cycle() will infinitely repeat the passed sequence:
>>> import itertools
>>> cs = itertools.cycle('ABC') # Note that strings are also a type of sequence
>>> for c in cs:
... print(c)
...
'A'
'B'
'C'
'A'
'B'
'C'
...
This also runs indefinitely.
repeat() is responsible for infinitely repeating a single element, but you can specify the number of repetitions by providing a second argument:
>>> ns = itertools.repeat('A', 3)
>>> for n in ns:
... print(n)
...
A
A
A
Infinite sequences only iterate infinitely when used in a for loop. If you only create an iterator object, it will not generate an infinite number of elements in advance—in fact, it is impossible to create an infinite number of elements in memory.
Although infinite sequences can be iterated indefinitely, we usually extract a finite sequence using conditional judgment functions such as takewhile():
>>> natuals = itertools.count(1)
>>> ns = itertools.takewhile(lambda x: x <= 10, natuals)
>>> list(ns)
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
Several iterator manipulation functions provided by itertools are even more useful:
chain() can concatenate a group of iterable objects to form a larger iterator:
>>> for c in itertools.chain('ABC', 'XYZ'):
... print(c)
# Iteration result: 'A' 'B' 'C' 'X' 'Y' 'Z'
groupby() picks out adjacent duplicate elements from the iterator and groups them together:
>>> for key, group in itertools.groupby('AAABBBCCAAA'):
... print(key, list(group))
...
A ['A', 'A', 'A']
B ['B', 'B', 'B']
C ['C', 'C']
A ['A', 'A', 'A']
In fact, the selection rule is implemented through a function: as long as the return values of two elements passed to the function are equal, these two elements are considered to be in the same group, and the function’s return value serves as the key of the group. If we want to group while ignoring case, we can make both elements ‘A’ and ‘a’ return the same key:
>>> for key, group in itertools.groupby('AaaBBbcCAAa', lambda c: c.upper()):
... print(key, list(group))
...
A ['A', 'a', 'a']
B ['B', 'B', 'b']
C ['c', 'C']
A ['A', 'A', 'a']
The value of pi (π) can be calculated using the formula:
Using the itertools module provided by Python, 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, ...
# step 2: Take the first N terms of the sequence: 1, 3, 5, 7, 9, ..., 2*N-1.
# step 3: Add positive/negative signs and divide by 4: 4/1, -4/3, 4/5, -4/7, 4/9, ...
# step 4: Sum the terms:
return 3.14
# Test:
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')
The itertools module provides only functions for handling iteration functionality. Their return values are not lists, but Iterators—calculations are only performed when iterating with a for loop.
use_itertools.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import itertools
natuals = itertools.count(1)
for n in natuals:
print(n)
if n >= 100:
break
cs = itertools.cycle("ABC")
t = 10
for c in cs:
print(c)
t = t - 1
if t == 0:
break