Модуль itertools: инструменты итерации

Модуль itertools — коллекция быстрых и памятеэффективных инструментов для работы с итерируемыми объектами. Функции из itertools возвращают итераторы (не списки), поэтому не создают данные в памяти до момента потребления. Это важно при работе с большими последовательностями.

Счётчик и повторение

import itertools

# count(start, step) — бесконечная арифметическая прогрессия:
for n in itertools.count(1, 2):
    if n > 10:
        break
    print(n, end=" ")
# 1 3 5 7 9

# repeat(value, times) — повторить значение N раз:
result = list(itertools.repeat("X", 5))
print(result)   # ['X', 'X', 'X', 'X', 'X']

# cycle(iterable) — бесконечное циклическое повторение:
colors = itertools.cycle(["R", "G", "B"])
for _ in range(7):
    print(next(colors), end=" ")
# R G B R G B R

chain — объединение последовательностей

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

import itertools

a = [1, 2, 3]
b = [4, 5]
c = [6, 7, 8, 9]

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

# chain.from_iterable — если у вас список списков:
nested = [[1, 2], [3, 4], [5, 6]]
flat = list(itertools.chain.from_iterable(nested))
print(flat)   # [1, 2, 3, 4, 5, 6]

islice — срез итератора

islice работает как срез списка [start:stop:step], но для любого итератора:

import itertools

numbers = range(100)
first_five = list(itertools.islice(numbers, 5))
print(first_five)   # [0, 1, 2, 3, 4]

# islice(iterable, start, stop, step):
every_third = list(itertools.islice(range(20), 0, 20, 3))
print(every_third)  # [0, 3, 6, 9, 12, 15, 18]

product и permutations

product — декартово произведение (все комбинации), аналог вложенных циклов:

import itertools

# Все возможные пары:
pairs = list(itertools.product([1, 2], ["a", "b"]))
print(pairs)
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]

# Эквивалент вложенного for:
# [(x, y) for x in [1, 2] for y in ["a", "b"]]

combinations и permutations — комбинаторика:

import itertools

items = [1, 2, 3]

# Перестановки (порядок важен):
for p in itertools.permutations(items, 2):
    print(p, end=" ")
# (1, 2) (1, 3) (2, 1) (2, 3) (3, 1) (3, 2)

print()

# Комбинации (порядок неважен):
for c in itertools.combinations(items, 2):
    print(c, end=" ")
# (1, 2) (1, 3) (2, 3)

groupby — группировка

groupby группирует последовательные элементы с одинаковым ключом (данные должны быть отсортированы):

import itertools

data = [
    {"name": "Алиса", "dept": "HR"},
    {"name": "Борис", "dept": "HR"},
    {"name": "Вера",  "dept": "IT"},
    {"name": "Глеб",  "dept": "IT"},
]

data.sort(key=lambda x: x["dept"])
for dept, group in itertools.groupby(data, key=lambda x: x["dept"]):
    names = [p["name"] for p in group]
    print(f"{dept}: {names}")
# HR: ['Алиса', 'Борис']
# IT: ['Вера', 'Глеб']

Проверь себя

Что вернёт list(itertools.chain([1, 2], [3], [4, 5]))?

chain объединяет все элементы по порядку: [1, 2, 3, 4, 5].

Итог

chain(*iterables) — последовательное объединение; chain.from_iterable — из списка списков.
islice(it, stop) — взять первые N элементов из итератора.
product(a, b) — декартово произведение; combinations / permutations — комбинаторика.
groupby(it, key) — группировка соседних элементов по ключу (нужна предварительная сортировка).
Все функции возвращают итераторы — эффективны по памяти.

Модуль 9 завершён. В следующем модуле перейдём к файлам: чтение, запись, работа с путями, CSV и JSON.

Модуль itertools: инструменты итерации