Python itertools
Table of Contents
1. Background
The itertools module provides memory-efficient tools for working with
iterators. These functions create iterators for efficient looping.
2. Infinite Iterators
from itertools import count, cycle, repeat
# count(start, step) - infinite counter
for i in count(10, 2):
if i > 20:
break
print(i) # 10, 12, 14, 16, 18, 20
# cycle(iterable) - repeat indefinitely
colors = cycle(['red', 'green', 'blue'])
[next(colors) for _ in range(5)]
# ['red', 'green', 'blue', 'red', 'green']
# repeat(element, times) - repeat element
list(repeat('hello', 3)) # ['hello', 'hello', 'hello']
3. Combinatoric Iterators
from itertools import permutations, combinations, product
# permutations - all orderings
list(permutations('ABC', 2))
# [('A','B'), ('A','C'), ('B','A'), ('B','C'), ('C','A'), ('C','B')]
# combinations - unique selections
list(combinations('ABC', 2))
# [('A','B'), ('A','C'), ('B','C')]
# product - cartesian product
list(product('AB', '12'))
# [('A','1'), ('A','2'), ('B','1'), ('B','2')]
4. Terminating Iterators
from itertools import chain, islice, groupby, takewhile, dropwhile
# chain - combine iterables
list(chain([1, 2], [3, 4], [5])) # [1, 2, 3, 4, 5]
# islice - slice an iterator
list(islice(range(100), 5, 10)) # [5, 6, 7, 8, 9]
# groupby - group consecutive elements
data = [('a', 1), ('a', 2), ('b', 3), ('b', 4)]
for key, group in groupby(data, key=lambda x: x[0]):
print(key, list(group))
# a [('a', 1), ('a', 2)]
# b [('b', 3), ('b', 4)]
# takewhile/dropwhile
list(takewhile(lambda x: x < 5, [1, 3, 5, 2, 1])) # [1, 3]
list(dropwhile(lambda x: x < 5, [1, 3, 5, 2, 1])) # [5, 2, 1]
5. Accumulation
from itertools import accumulate
import operator
# Running sum
list(accumulate([1, 2, 3, 4])) # [1, 3, 6, 10]
# Running product
list(accumulate([1, 2, 3, 4], operator.mul)) # [1, 2, 6, 24]
# Running max
list(accumulate([3, 1, 4, 1, 5], max)) # [3, 3, 4, 4, 5]
6. Practical Examples
from itertools import zip_longest, starmap, filterfalse
# zip_longest - zip with fill value
list(zip_longest([1, 2], [3, 4, 5], fillvalue=0))
# [(1, 3), (2, 4), (0, 5)]
# starmap - unpack arguments
list(starmap(pow, [(2, 3), (3, 2), (10, 2)])) # [8, 9, 100]
# filterfalse - inverse of filter
list(filterfalse(lambda x: x % 2, range(10))) # [0, 2, 4, 6, 8]