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

В декабре 2019 года мы подробно рассказали о модуле collections. Другой важный компонент стандартной библиотеки – itertiools.

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

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

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

             import itertools          

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

Если вы владеете Jupyter Notebook, блокнот этой статьи доступен на GitHub.

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

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

             >>> cnt = itertools.count(start=2020, step=4) >>> next(cnt) 2020 >>> next(cnt) 2024 >>> next(cnt) 2028          

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

             >>> days = [366]*4 >>> list(zip(itertools.count(2020, 4), days)) [(2020, 366), (2024, 366), (2028, 366), (2032, 366)]          

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

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

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

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

             >>> list(zip(range(0, 10), range(0, 5))) [(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)          

             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):

             >>> list(itertools.accumulate(range(1, 10))) [1, 3, 6, 10, 15, 21, 28, 36, 45]          

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

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

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

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

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

             >>> waltz = itertools.cycle(['и раз', 'и два', 'и три']) >>> next(waltz) 'и раз' >>> next(waltz) 'и два' >>> next(waltz) 'и три' >>> next(waltz) 'и раз'          

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

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

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

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

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

Все числа последовательности nums возведены в степень 2. Итератор вызывается столько раз, сколько необходимо, не нужно думать о числе элементов в последовательности.

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

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

             >>> squares = itertools.starmap(pow, [(0, 2), (1, 2), (2, 2)]) >>> list(squares) [0, 1, 4]          

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

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

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

Например, мы хотим составить трёхцветный флаг из лент цветных тканей. Есть четыре цвета лент. Все варианты выбора тканей без учёта их расположения:

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

             ('белый', 'жёлтый', 'синий') ('белый', 'жёлтый', 'красный') ('белый', 'синий', 'красный') ('жёлтый', 'синий', 'красный')          

Порядок следования не имеет значения, поэтому все тройки цветов уникальны.

8. Комбинаторика: перестановки

Перестановки – те же сочетания, для которых важен порядок следования элементов. В продолжение предыдущего примера определим все варианты как мы можем составить флаг с учётом порядка следования цветов:

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

             ('белый', 'жёлтый', 'синий') ('белый', 'жёлтый', 'красный') ('белый', 'синий', 'жёлтый') ('белый', 'синий', 'красный') ('белый', 'красный', 'жёлтый') ('белый', 'красный', 'синий') ('жёлтый', 'белый', 'синий') ('жёлтый', 'белый', 'красный') ('жёлтый', 'синий', 'белый') ('жёлтый', 'синий', 'красный') ('жёлтый', 'красный', 'белый') ('жёлтый', 'красный', 'синий') ('синий', 'белый', 'жёлтый') ('синий', 'белый', 'красный') ('синий', 'жёлтый', 'белый') ('синий', 'жёлтый', 'красный') ('синий', 'красный', 'белый') ('синий', 'красный', 'жёлтый') ('красный', 'белый', 'жёлтый') ('красный', 'белый', 'синий') ('красный', 'жёлтый', 'белый') ('красный', 'жёлтый', 'синий') ('красный', 'синий', 'белый') ('красный', 'синий', 'жёлтый')          

Как и ожидалось, число комбинаций с учетом порядка элементов гораздо больше.

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

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

Например, есть пин-код из четырех цифр. На каждой позиции стоит цифра от 0 до 9. Позиции не зависят друг от друга. Переберем все возможные коды:

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

             (0, 0, 0, 0) (0, 0, 0, 1) (0, 0, 0, 2) ... (9, 9, 9, 8) (9, 9, 9, 9)          

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

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

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

             ('A', 'A') ('A', 'B') ('A', 'C') ('A', 'D') ('B', 'B') ('B', 'C') ('B', 'D') ('C', 'C') ('C', 'D') ('D', 'D')          

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

Метод itertools.product() можно использовать не только для размещений с повторениями.

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

Например, найдём обозначения всех полей шахматной доски:

             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=' ')          

             a1 a2 a3 a4 a5 a6 a7 a8 b1 b2 b3 b4 b5 b6 b7 b8 c1 c2 c3 c4 c5 c6 c7 c8 d1 d2 d3 d4 d5 d6 d7 d8 e1 e2 e3 e4 e5 e6 e7 e8 f1 f2 f3 f4 f5 f6 f7 f8 g1 g2 g3 g4 g5 g6 g7 g8 h1 h2 h3 h4 h5 h6 h7 h8           

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

Иногда необходимо использовать нескольков итераторов. И независимо, и цепочкой один за другим. Для объединения итераторов используйте itertools.chain(*iterables).

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

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

С помощью itertools.chain() также можно добавлять отдельные элементы в начало итератора:

             >>> def prepend(value, iterator): ...    return itertools.chain([value], iterator) ... >>> list(prepend(1, [2, 3, 4])) [1, 2, 3, 4]          

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

Альтернативным конструктором itertools.chain() служит itertools.chain.from_iterable(). Метод принимает один итерируемый объект. Сравните их вызовы:

             >>> list(itertools.chain('ABC', 'DEF')) ['A', 'B', 'C', 'D', 'E', 'F'] >>> list(itertools.chain.from_iterable(['ABC', 'DEF'])) ['A', 'B', 'C', 'D', 'E', 'F']          

Последний конструктор удобно использовать для объединения списков:

             >>> list_of_lists = [[1, 2], [3, 4, 5], [6, 7, 8, 9]] >>> list(itertools.chain.from_iterable(list_of_lists)) [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 = '')          

             Строка 1 Строка 2 Строка 3          

Функция itertools.islice() позволяет итерироваться по любым объектам в формате среза. Например, следующая функция возвращает n первых элементов итерируемого объекта в виде списка:

             def take(n, iterable):     return list(islice(iterable, n))          

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

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

             >>> numbers = [0, 1, 2, 3, 2, 1, 0] >>> selectors = [True, True, False, True] >>> list(itertools.compress(numbers, selectors)) [0, 1, 3]          

Метод itertools.filterfalse() дополняет обычный фильтр filter():

             def filter_func(n):     if n < 2:         return True     return False  print(list(filter(filter_func, numbers))) print(list(itertools.filterfalse(filter_func, numbers)))          

             [0, 1, 1, 0] [2, 3, 2]          

16. Фильтрация до последнего истинного (или с первого ложного) элемента

Если необходимо отобрать объекты, стоящие после неудовлетворяющего условию элемента, используем itertools.dropwhile():

             >>> list(itertools.dropwhile(filter_func, numbers)) [2, 3, 2, 1, 0]          

Метод itertools.takewhile() наоборот выведет элементы, удовлетворяющие условию вплоть до объекта, прерывающего цепочку истинных элементов:

             >>> list(itertools.takewhile(filter_func, numbers)) [0, 1]          

17. Группировка по ключу

Инструмент itertools.groupby() объединяет смежные словари в группы по общему ключу. Например, сгруппируем студентов с одинаковой оценкой:

             people = [{"Имя": "Петр",            "Отчество": "Петрович",            "Фамилия": "Петров",            "Оценка":5},           {"Имя": "Ольга",            "Отчество": "Алексеевна",            "Фамилия": "Иванова",            "Оценка":5},           {"Имя": "Николай",            "Отчество": "Николаевич",            "Фамилия": "Николаев",            "Оценка":4},           {"Имя": "Федор",            "Отчество": "Владимирович",            "Фамилия": "Иванов",            "Оценка":3},           {"Имя": "Владимир",            "Отчество": "Федорович",            "Фамилия": "Иванов",            "Оценка":3}]  def get_mark(person):     return person['Оценка']  person_marks = itertools.groupby(people, get_mark)  for key, group in person_marks:     print(key)     for person in group:         print(person)     print()          

             5 {'Имя': 'Петр', 'Отчество': 'Петрович', 'Фамилия': 'Петров', 'Оценка': 5} {'Имя': 'Ольга', 'Отчество': 'Алексеевна', 'Фамилия': 'Иванова', 'Оценка': 5}  4 {'Имя': 'Николай', 'Отчество': 'Николаевич', 'Фамилия': 'Николаев', 'Оценка': 4}  3 {'Имя': 'Федор', 'Отчество': 'Владимирович', 'Фамилия': 'Иванов', 'Оценка': 3} {'Имя': 'Владимир', 'Отчество': 'Федорович', 'Фамилия': 'Иванов', 'Оценка': 3}          

Обратите внимание, что группировка работает только со смежными объектами. Предварительно отсортируйте данные.

18. Репликация итераторов

Функция itertools.tee() создаёт из одного итерируемого объекта два итератора:

             >>> letters = 'abc' >>> it1, it2 = itertools.tee(letters) >>> next(it1) 'a' >>> next(it1) 'b' >>> next(it2) 'a'          

Эти итераторы соответствуют одной последовательности, но независимы друг от друга.

19. Повторение последовательности заданное число раз

Волшебная сила itertools – в умении комбинировать итераторы, чтобы писать быстрый, эффективный и ясный код.

Например, сочетание itertools.chain() и itertools.from_iterable() даёт ограниченный вариант бесконечного itertools.cycle():

             def ncycles(iterable, n):     return itertools.chain(itertools.from_iterable(repeat(tuple(iterable), n)))          

20. Уникальные элементы последовательности

Напоследок напишем функцию, выводящую все уникальные элементы в том же порядке, как они появлялись в исходной последовательности. Опциональный аргумент – ссылка на функцию определения эквивалентов, представленных по-разному. Например, заглавных и строчных букв.

             def unique_everseen(iterable, key=None):     seen = set()     seen_add = seen.add     if key is None:         for element in itertools.filterfalse(seen.__contains__, iterable):             seen_add(element)             yield element     else:         for element in iterable:             k = key(element)             if k not in seen:                 seen_add(k)                 yield element          

             >>> list(unique_everseen('Абракадааааабра')) ['А', 'б', 'р', 'а', 'к', 'д'] >>> list(unique_everseen('Абракадааааабра', str.lower)) ['А', 'б', 'р', 'к', 'д']          

Заключение

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

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

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

             pip install more-itertools          

Интересны ли вам такие обзоры инструментов Python? Будем рады узнать в комментариях.