Модуль itertools в Python

Модуль itertools из стандартной библиотеки Python — это набор быстрых и экономичных по памяти строительных блоков для работы с итераторами. Каждая функция в itertools возвращает итератор — он генерирует значения по требованию, а не строит список в памяти — что делает модуль идеальным для больших наборов данных, бесконечных последовательностей и составных конвейеров обработки данных.

В этой главе рассматриваются все три категории функций itertools: бесконечные итераторы (count, cycle, repeat), комбинаторные итераторы (product, permutations, combinations, combinations_with_replacement) и завершающие итераторы (chain, islice, groupby, compress, filterfalse, takewhile, dropwhile, starmap, zip_longest, accumulate, pairwise).

Установка не требуется — itertools входит в стандартный дистрибутив Python 3:

import itertools

Зачем нужен itertools?

Рассмотрим задачу получения первых 10 кратных некоторого числа. Без itertools потребуется список или счётчик вручную. С itertools.count и itertools.islice намерение сразу очевидно, а потребление памяти остаётся постоянным:

import itertools

multiples = itertools.islice(itertools.count(0, 7), 10)
print(list(multiples))
# [0, 7, 14, 21, 28, 35, 42, 49, 56, 63]

Философия itertools: создать маленький, корректный блок, затем скомпоновать его с другими. Соединение двух функций itertools работает быстрее и менее подвержено ошибкам, чем написание эквивалентного цикла вручную.

Бесконечные итераторы

Эти итераторы производят значения бесконечно. Всегда сочетайте их с islice, конструкцией for … break или другим ограничивающим механизмом, чтобы избежать бесконечного цикла.

count(start=0, step=1)

count производит равномерно распределённую последовательность чисел. По сути это range без верхней границы с поддержкой чисел с плавающей точкой и отрицательного шага.

import itertools

# Integer counter
for n in itertools.islice(itertools.count(10), 5):
    print(n, end=' ')
# 10 11 12 13 14

print()

# Float step
for n in itertools.islice(itertools.count(0.0, 0.5), 5):
    print(n, end=' ')
# 0.0 0.5 1.0 1.5 2.0

print()

# Countdown
for n in itertools.islice(itertools.count(100, -10), 5):
    print(n, end=' ')
# 100 90 80 70 60

count удобен, когда нужно пронумеровать элементы итерируемого объекта, не зная заранее их количество, — это идиома enumerate, но с пользовательским начальным значением и шагом.

cycle(iterable)

cycle бесконечно повторяет элементы любого итерируемого объекта.

import itertools

colours = itertools.cycle(['red', 'green', 'blue'])
for i, colour in enumerate(colours):
    if i == 7:
        break
    print(colour, end=' ')
# red green blue red green blue red

Практическое применение — распределение участников по командам по принципу карусели:

import itertools

teams = itertools.cycle(['Alpha', 'Beta', 'Gamma'])
players = ['Alice', 'Bob', 'Carol', 'Dave', 'Eve']
assignments = {player: team for player, team in zip(players, teams)}
print(assignments)
# {'Alice': 'Alpha', 'Bob': 'Beta', 'Carol': 'Gamma', 'Dave': 'Alpha', 'Eve': 'Beta'}

repeat(object, times=None)

repeat возвращает один и тот же объект times раз (или бесконечно, если times не задан).

import itertools

# Finite repeat
print(list(itertools.repeat('hello', 3)))
# ['hello', 'hello', 'hello']

# Used as a fixed argument supplier in map()
squares = list(map(pow, range(1, 6), itertools.repeat(2)))
print(squares)
# [1, 4, 9, 16, 25]

Паттерн map(pow, range(1, 6), repeat(2)) — распространённая идиома для передачи константного второго аргумента в двухаргументную функцию.

Комбинаторные итераторы

Эти итераторы производят все комбинации, перестановки или декартовы произведения входного итерируемого объекта. Они необходимы для перебора с возвратом, генерации тестовых случаев и комбинаторных задач.

product(*iterables, repeat=1)

product вычисляет декартово произведение — все упорядоченные комбинации, где один элемент берётся из каждого итерируемого объекта. Это эквивалент вложенных циклов for.

import itertools

suits = ['Hearts', 'Diamonds']
ranks = ['A', 'K', 'Q']
deck = list(itertools.product(suits, ranks))
print(deck)
# [('Hearts', 'A'), ('Hearts', 'K'), ('Hearts', 'Q'),
#  ('Diamonds', 'A'), ('Diamonds', 'K'), ('Diamonds', 'Q')]

Используйте repeat для вычисления произведения итерируемого объекта с самим собой несколько раз:

import itertools

# All 2-digit binary numbers
binary_pairs = list(itertools.product([0, 1], repeat=2))
print(binary_pairs)
# [(0, 0), (0, 1), (1, 0), (1, 1)]

Важное замечание: product загружает входные итерируемые объекты в память (чтобы обеспечить несколько проходов), поэтому не передавайте в него огромные итераторы.

permutations(iterable, r=None)

permutations производит все упорядоченные расстановки из r элементов, взятых из входного объекта. Если r не указан, используются все элементы.

import itertools

# All orderings of 3 letters
perms = list(itertools.permutations('ABC'))
print(perms)
# [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'),
#  ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]
print(len(perms))   # 6  (3! = 6)

# 2-element permutations
perms2 = list(itertools.permutations('ABC', 2))
print(perms2)
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
print(len(perms2))  # 6  (3 * 2 = 6)

В перестановках порядок важен — ('A', 'B') и ('B', 'A') являются разными результатами.

combinations(iterable, r)

combinations производит все неупорядоченные выборки из r элементов. В отличие от permutations, порядок не важен — каждое подмножество встречается только один раз.

import itertools

# All 2-element subsets of [1, 2, 3, 4]
combos = list(itertools.combinations([1, 2, 3, 4], 2))
print(combos)
# [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
print(len(combos))  # 6  (C(4,2) = 6)

Распространённый сценарий использования — проверка всех пар элементов на некоторое свойство:

import itertools

words = ['bat', 'tab', 'cat', 'tac']
anagram_pairs = [
    (a, b) for a, b in itertools.combinations(words, 2)
    if sorted(a) == sorted(b)
]
print(anagram_pairs)
# [('bat', 'tab'), ('cat', 'tac')]

combinations_with_replacement(iterable, r)

Аналог combinations, но позволяет каждому элементу встречаться в выборке более одного раза.

import itertools

# All 2-element combinations with repetition from [1, 2, 3]
combos = list(itertools.combinations_with_replacement([1, 2, 3], 2))
print(combos)
# [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

Это полезно для генерации всех возможных бросков кубика, последовательностей подбрасывания монеты или вариантов выбора символов для пароля.

Комбинаторные функции: краткое сравнение

Завершающие итераторы

Завершающие итераторы обрабатывают конечный входной объект и останавливаются, когда он исчерпан.

chain(*iterables)

chain обрабатывает несколько итерируемых объектов как единую непрерывную последовательность без создания нового списка.

import itertools

a = [1, 2, 3]
b = (4, 5)
c = range(6, 9)
combined = list(itertools.chain(a, b, c))
print(combined)
# [1, 2, 3, 4, 5, 6, 7, 8]

chain.from_iterable принимает единственный итерируемый объект итерируемых объектов — удобно, когда число последовательностей неизвестно заранее:

import itertools

nested = [[1, 2], [3, 4], [5, 6]]
flat = list(itertools.chain.from_iterable(nested))
print(flat)
# [1, 2, 3, 4, 5, 6]

Это быстрая, экономичная по памяти альтернатива выражению [item for sublist in nested for item in sublist].

islice(iterable, stop) / islice(iterable, start, stop, step=1)

islice нарезает любой итератор — включая бесконечные — без загрузки его в память. Аргументы аналогичны нотации slice в Python, но принимают только неотрицательные целые числа.

import itertools

# First 5 elements
print(list(itertools.islice(range(100), 5)))
# [0, 1, 2, 3, 4]

# Elements 10–14 (start inclusive, stop exclusive)
print(list(itertools.islice(range(100), 10, 15)))
# [10, 11, 12, 13, 14]

# Every other element from position 0 to 10
print(list(itertools.islice(range(20), 0, 10, 2)))
# [0, 2, 4, 6, 8]

islice не поддерживает отрицательные индексы или отрицательный шаг (в отличие от стандартного среза списка).

groupby(iterable, key=None)

groupby группирует последовательно идущие элементы с одинаковым значением ключа. Возвращает пары (ключ, итератор_группы).

import itertools

data = [
    ('fruit', 'apple'),
    ('fruit', 'banana'),
    ('veggie', 'carrot'),
    ('veggie', 'broccoli'),
    ('fruit', 'cherry'),
]

for category, group in itertools.groupby(data, key=lambda x: x[0]):
    items = [item[1] for item in group]
    print(f'{category}: {items}')
# fruit: ['apple', 'banana']
# veggie: ['carrot', 'broccoli']
# fruit: ['cherry']

Важное замечание: groupby группирует только последовательно идущие одинаковые элементы. Если данные не отсортированы по ключу, похожие элементы в разных позициях образуют отдельные группы (как показано выше — 'cherry' начинает новую группу 'fruit', а не попадает в первую). Всегда сортируйте данные по ключу перед вызовом groupby:

import itertools

data = [
    ('fruit', 'apple'),
    ('veggie', 'carrot'),
    ('fruit', 'banana'),
    ('veggie', 'broccoli'),
    ('fruit', 'cherry'),
]

# Sort first, then group
sorted_data = sorted(data, key=lambda x: x[0])
for category, group in itertools.groupby(sorted_data, key=lambda x: x[0]):
    items = [item[1] for item in group]
    print(f'{category}: {items}')
# fruit: ['apple', 'banana', 'cherry']
# veggie: ['carrot', 'broccoli']

Также учтите, что итератор группы становится недействительным после перехода к следующему ключу — исчерпывайте каждую группу до вызова next() на внешнем итераторе.

compress(data, selectors)

compress фильтрует data, оставляя только те элементы, у которых соответствующее значение selector истинно.

import itertools

names = ['Alice', 'Bob', 'Carol', 'Dave', 'Eve']
active = [True, False, True, True, False]

result = list(itertools.compress(names, active))
print(result)
# ['Alice', 'Carol', 'Dave']

compress эквивалентен выражению [d for d, s in zip(data, selectors) if s], но работает быстрее и без промежуточного списка.

filterfalse(predicate, iterable)

filterfalse — дополнение встроенной функции filter: возвращает элементы, для которых предикат возвращает False.

import itertools

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

# Keep only odd numbers (those that fail the even test)
odds = list(itertools.filterfalse(lambda x: x % 2 == 0, numbers))
print(odds)
# [1, 3, 5, 7, 9]

takewhile(predicate, iterable)

takewhile возвращает элементы до тех пор, пока предикат равен True, затем немедленно останавливается — даже если последующие элементы удовлетворяют предикату.

import itertools

data = [2, 4, 6, 3, 8, 10]

# Stop as soon as an odd number appears
evens_from_start = list(itertools.takewhile(lambda x: x % 2 == 0, data))
print(evens_from_start)
# [2, 4, 6]

dropwhile(predicate, iterable)

dropwhile — зеркало takewhile: пропускает элементы, пока предикат равен True, затем возвращает все оставшиеся элементы (включая те, для которых предикат снова стал бы True).

import itertools

data = [2, 4, 6, 3, 8, 10]

# Drop leading even numbers, yield everything from the first odd onward
result = list(itertools.dropwhile(lambda x: x % 2 == 0, data))
print(result)
# [3, 8, 10]

takewhile и dropwhile удобны при обработке лог-файлов или потоков данных, когда нужно пропустить заголовочный раздел или остановиться на строке-маркере.

starmap(function, iterable)

starmap применяет функцию к каждому элементу итерируемого объекта, распаковывая элемент как позиционные аргументы. Это аналог map, но для итерируемых объектов из кортежей.

import itertools

pairs = [(2, 3), (4, 2), (10, 3)]
results = list(itertools.starmap(pow, pairs))
print(results)
# [8, 16, 1000]

Сравните с map(pow, [2, 4, 10], [3, 2, 3]) — starmap работает, когда аргументы уже упакованы в кортежи.

zip_longest(*iterables, fillvalue=None)

Встроенный zip останавливается на самом коротком итерируемом объекте. zip_longest дополняет более короткие объекты значением fillvalue, чтобы все они были обработаны полностью.

import itertools

a = [1, 2, 3]
b = ['a', 'b', 'c', 'd', 'e']

print(list(zip(a, b)))
# [(1, 'a'), (2, 'b'), (3, 'c')]  — b's 'd' and 'e' are lost

print(list(itertools.zip_longest(a, b, fillvalue=0)))
# [(1, 'a'), (2, 'b'), (3, 'c'), (0, 'd'), (0, 'e')]

accumulate(iterable, func=operator.add, *, initial=None)

accumulate вычисляет накопленные суммы (или любую другую нарастающую агрегацию). По умолчанию суммирует, но можно передать любую двухаргументную функцию.

import itertools
import operator

numbers = [1, 2, 3, 4, 5]

# Running sum (default)
print(list(itertools.accumulate(numbers)))
# [1, 3, 6, 10, 15]

# Running product
print(list(itertools.accumulate(numbers, operator.mul)))
# [1, 2, 6, 24, 120]

# Running maximum
data = [3, 1, 4, 1, 5, 9, 2, 6]
print(list(itertools.accumulate(data, max)))
# [3, 3, 4, 4, 5, 9, 9, 9]

Параметр initial (Python 3.8+) добавляет начальное значение перед первым элементом:

import itertools

print(list(itertools.accumulate([1, 2, 3], initial=100)))
# [100, 101, 103, 106]

pairwise(iterable)

pairwise (Python 3.10+) возвращает последовательные перекрывающиеся пары из итерируемого объекта.

import itertools

data = [1, 2, 3, 4, 5]
print(list(itertools.pairwise(data)))
# [(1, 2), (2, 3), (3, 4), (4, 5)]

Это полезно для вычисления разностей между последовательными значениями или для логики скользящего окна с размером ровно 2:

import itertools

prices = [10.0, 12.5, 11.0, 13.5, 15.0]
changes = [b - a for a, b in itertools.pairwise(prices)]
print(changes)
# [2.5, -1.5, 2.5, 1.5]

До Python 3.10 эквивалентом был zip(data, data[1:]) (работает для последовательностей) или подход на основе tee (работает для произвольных итераторов).

Составление конвейеров из itertools

Настоящая мощь itertools проявляется при комбинировании функций. Так как каждая функция возвращает итератор, их можно соединять в цепочки без промежуточных списков.

Пример: топ-3 слов по частоте в тексте

import itertools
import operator

text = "the quick brown fox jumps over the lazy dog the fox"
words = text.split()

# Sort words so groupby can collect identical words together
sorted_words = sorted(words)

# Count each word using groupby
word_counts = (
    (key, sum(1 for _ in group))
    for key, group in itertools.groupby(sorted_words)
)

# Sort by count descending, take the top 3
top3 = list(itertools.islice(
    sorted(word_counts, key=operator.itemgetter(1), reverse=True),
    3
))
print(top3)
# [('the', 3), ('fox', 2), ('brown', 1)]

Пример: разбивка итерируемого объекта на фрагменты фиксированного размера

import itertools

def batched(iterable, n):
    """Yield successive n-sized tuples from iterable."""
    it = iter(iterable)
    while chunk := tuple(itertools.islice(it, n)):
        yield chunk

data = range(10)
for batch in batched(data, 3):
    print(batch)
# (0, 1, 2)
# (3, 4, 5)
# (6, 7, 8)
# (9,)

Python 3.12 включает встроенный itertools.batched, поэтому на современном Python можно заменить вспомогательную функцию выше на itertools.batched(data, 3).

Краткий справочник

По концепциям ленивых вычислений, лежащим в основе itertools, см. Генераторы Python и Итераторы Python. По вспомогательным инструментам в функциональном стиле, дополняющим itertools, см. Лямбда-функции Python и Модуль collections в Python.

Практика