Skip to content

Итерируемся правильно itertools в Python

Модуль itertools содержит строительные блоки итераторов, основанные на конструкциях из языков программирования APL, Haskell и SML. Ниже мы опишем набор быстрых и эффективных в отношении памяти инструментов, полезных как самостоятельно, так и в сочетании. Вместе они образуют «алгебру итераторов» для программ на чистом Python.

Цель публикации – в сжатой форме рассмотреть распространённые примеры и шаблоны использования модуля itertools.

Начнем с импорта:

import itertools

Чтобы лучше запомнить функции модуля, мы не станем использовать конструкцию from ... import *, а будем обращаться к методам модуля через его имя.

Если вы владеете Jupyter Notebook, блокнот этой статьи доступен на GitHub. Соответственно код легко запустить в интерактивном режиме с помощью Colab.

1. Бесконечный счётчик

Функция itertools.count(start=0, step=1) создаёт бесконечный итератор. Можно задать начальное значение и шаг итерирования.

cnt = itertools.count(start=2020, step=4)
print(next(cnt))  # Output: 2020
print(next(cnt))  # Output: 2024
print(next(cnt))  # Output: 2028

Пример использования итератора в zip-функции:

days = [366] * 4
print(list(zip(itertools.count(2020, 4), days)))
# Output: [(2020, 366), (2024, 366), (2028, 366), (2032, 366)]

Чтобы продолжить счёт при прерывании выполнения программы, передайте последнее значение новому объекту итератора в виде параметра start.

Если необходимо подсчитывать число вхождений элементов в список или кортеж, обратите внимание на Counter() из модуля collections.

2. Упаковка по более длинной последовательности

Если последовательности имеют неодинаковую длину, zip() ограничивается самой короткой:

print(list(zip(range(0, 10), range(0, 5))))
# Output: [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]

Но такое сокращение может быть неудобно из-за потери информации. Чтобы сохранить обе последовательности, используйте itertools.zip_longest():

for (i, j) in itertools.zip_longest(range(0, 10), range(0, 5)):
    print(i, j)
# Output:
# 0 0
# 1 1
# 2 2
# 3 3
# 4 4
# 5 None
# 6 None
# 7 None
# 8 None
# 9 None

Вместо None функция может подставлять значение, переданное аргументу fillvalue.

3. Аккумулирующий итератор

Суммирование нарастающим (накопительным) итогом – вид сложения последовательности чисел. Например, так считается квартальная прибыль. Каждый элемент складывается с суммой всех предшествовавших элементов. В следующем примере 1 и 2 даёт 3, сумма 1, 2, 3 равна 6 и т. д. Описанный тип работы с последовательностью воплощен в itertools.accumulate(iterable, func=operator.add, *, initial=None):

print(list(itertools.accumulate(range(1, 10))))
# Output: [1, 3, 6, 10, 15, 21, 28, 36, 45]

По умолчанию к элементам применяется operator.add. Можно, например, указать оператор умножения:

import operator
print(list(itertools.accumulate(range(1, 10), operator.mul)))
# Output: [1, 2, 6, 24, 120, 720, 5040, 40320, 362880]

4. Бесконечный итератор последовательности

С помощью itertools.cycle() создаётся кольцевой итератор. Прийдя к последнему значению, он вновь начинает с первого:

waltz = itertools.cycle(['и раз', 'и два', 'и три'])
print(next(waltz))  # Output: 'и раз'
print(next(waltz))  # Output: 'и два'
print(next(waltz))  # Output: 'и три'
print(next(waltz))  # Output: 'и раз'

5. Бесконечный итератор одного объекта

Итератор, создаваемый itertools.repeat(), это вырожденный случай itertools.cycle(). Вместо последовательности повторяется одно и то же значение. Бесконечно или times раз:

s = "Птица Говорун отличается умом и сообразительностью"
rep = itertools.repeat(s, times=2)
print(next(rep))  # Output: 'Птица Говорун отличается умом и сообразительностью'
print(next(rep))  # Output: 'Птица Говорун отличается умом и сообразительностью'

Классический пример использования itertools.repeat() – итератор для map():

nums = range(10)
squares = map(pow, nums, itertools.repeat(2))
print(list(squares))  # Output: [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

6. Мапирование с распаковкой

Раз мы заговорили о map(), полезно рассказать и о itertools.starmap(). Этот метод принимает функцию и список кортежей аргументов. Как если бы использовался оператор *, отсюда и название:

squares = itertools.starmap(pow, [(0, 2), (1, 2), (2, 2)])
print(list(squares))  # Output: [0, 1, 4]

7. Комбинаторика: сочетания

Модуль itertools позволяет решать программные задачи, построенные на структурах комбинаторики.

Сочетания – выбранные из множества n объектов комбинации m объектов, отличающиеся хотя бы одним объектом. Порядок элементов не важен.

colors = ['белый', 'жёлтый', 'синий', 'красный']
for item in itertools.combinations(colors, 3):
    print(item)
# Output:
# ('белый', 'жёлтый', 'синий')
# ('белый', 'жёлтый', 'красный')
# ('белый', 'синий', 'красный')
# ('жёлтый', 'синий', 'красный')

8. Комбинаторика: размещения

Размещения – те же сочетания, для которых важен порядок следования элементов.

for item in itertools.permutations(colors, 3):
    print(item)

9. Комбинаторика: размещение с повторениями

Размещение с повторениями (выборка с возвращением) – это комбинаторное размещение объектов, в котором каждый объект может участвовать в размещении несколько раз.

digits = range(10)
pincode_vars = itertools.product(digits, repeat=4)
for var in pincode_vars:
    print(var)

10. Комбинаторика: размещение

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

letters = 'ABCD'
code_vars = itertools.combinations_with_replacement(letters, 2)
for var in code_vars:
    print(var)

11. Декартово произведение множеств

Декартово (прямое) произведение – множество, элементами которого являются все возможные упорядоченные пары элементов исходных множеств.

import string
letters = list(string.ascii_lowercase[:8])
digits = range(1, 9)
for (letter, digit) in itertools.product(letters, digits):
    print(letter + str(digit), end=' ')

12. Цепочки итераторов

Для объединения итераторов используйте itertools.chain(*iterables):

num_cards = [str(i) for i in range(2, 11)]
face_cards = ['В', 'Д', 'К', 'Т']
print(list(itertools.chain(num_cards, face_cards)))
# Output: ['2', '3', '4', '5', '6', '7', '8', '9', '10', 'В', 'Д', 'К', 'Т']

13. Плоский список из вложенного

Альтернативным конструктором itertools.chain() служит itertools.chain.from_iterable():

list_of_lists = [[1, 2], [3, 4, 5], [6, 7, 8, 9]]
print(list(itertools.chain.from_iterable(list_of_lists)))
# Output: [1, 2, 3, 4, 5, 6, 7, 8, 9]

14. Итератор среза

Срез – удобный инструмент списков, который доступен и для итераторов с помощью itertools.islice():

with open('test.txt', 'r') as f:
    header = itertools.islice(f, 3)
    for line in header:
        print(line, end='')

15. Фильтрация группы элементов

Функция compress() оставляет из итерируемых данных только те, что соответствуют позициям булевых селекторов:

numbers = [0, 1, 2, 3, 2, 1, 0]
selectors = [True, True, False, True]
print(list(itertools.compress(numbers, selectors)))
# Output: [0, 1, 3]

Заключение

Описывая приёмы использования itertools, мы попутно определили основные функции модуля.

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

Модуль itertools обеспечивает ключевые структуры итераторов Python. Другие шаблоны вы найдёте в специальной библиотеке примеров more-itertools.