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

Рассказываем с примерами кода о функциях модуля itertools – инструмента стандартной библиотеки Python, содержащего распространённые шаблоны итераторов. Бесконечные счётчики, сочетания и перестановки, итераторы среза и многое другое.

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

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

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

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

 import itertools 

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

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

Функция 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.

Если последовательности имеют неодинаковую длину, 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.

Суммирование нарастающим (накопительным) итогом – вид сложения последовательности чисел. Например, так считается квартальная прибыль Каждый элемент складывается с суммой всех предшествовавших элементов. В следующем примере 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] 

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

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

Итератор, создаваемый 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. Итератор вызывается столько раз, сколько необходимо, не нужно думать о числе элементов в последовательности.

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

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

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

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

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

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

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

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

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

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

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

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

Например, есть пин-код из четырех цифр. На каждой позиции стоит цифра от 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) 

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

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

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

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

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

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

 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  

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

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

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

 >>> 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] 

Альтернативным конструктором 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] 

Срез – удобный инструмент списков, который доступен и для итераторов с помощью 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)) 

Функция 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] 

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

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

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

 >>> list(itertools.takewhile(filter_func, numbers))
[0, 1] 
Итерируем правильно: 20 приемов использования в Python модуля itertools

Инструмент 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} 

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

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

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

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

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

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

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

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

 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? Будем рады узнать в комментариях.

proglib.io

Добавить комментарий

Ваш e-mail не будет опубликован.

пять × два =