Сообщить об ошибке.

Введение в модуль itertools Python

Python предоставляет отличный модуль itertools для создания собственных итераторов. Инструменты, предоставляемые itertools, работают быстро и эффективно используют память. Эти строительные блоки можно использовать для создания собственных специализированных итераторов, которые отлично впишутся в эффективный цикл for/in.

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

Содержание:


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

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

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


Итератор itertools.count() будет возвращать равномерно распределенные значения, начиная с числа, которое передается в качестве начального параметра start; этот итератор также принимает параметр шага step.

Простой пример:

>>> from itertools import count
>>> for i in count(10):
...     if i > 15: 
...         break
...     else:
...         print(i)

# 10
# 11
# 12
# 13
# 14
# 15

Здесь импортируется функция count из модуля itertools и создается цикл for/in. Для выхода из бесконечного цикла добавляется условная проверка, если итератор превысит число 15, в противном случае он распечатает, где вы находитесь в итераторе. Обратите внимание, что вывод начинается с 10, так как это начальное значение start.

Другой способ ограничить вывод этого бесконечного итератора - использовать другую функцию из itertools с именем itertools.islice().

Пример:

>>> from itertools import islice, count
>>> for i in islice(count(10), 5):
...     print(i)

# 10
# 11
# 12
# 13
# 14

В этом примере функция itertools.islice() выполняет срез бесконечного итератора (обычная операция среза последовательности не работает с типом итератор) и ограничивая его пятью элементами который затем перебирается в цикле for.


Итератор itertools.cycle() позволяет создать итератор, который будет бесконечно циклически перебирать серию значений. Передадим ему строку из 3 букв и посмотрим, что произойдет:

>>> from itertools import cycle
>>> for i, item in enumerate(cycle('XYZ')):
...     if i > 5:
...         break
...     print(item)

# X
# Y
# Z
# X
# Y

Здесь создается цикл for/in для бесконечного цикла трех букв: X, Y, Z. Для выхода из бесконечного цикла используется условие, которое проверяет счетчик, который генерирует встроенная функция enumerate.

Также можно использовать встроенную функцию next() для перемещения по итераторам, созданных при помощи модуля itertools.

>>> from itertools import cycle
>>> polys = ['triangle', 'square', 'pentagon']
>>> iterator = cycle(polys)
>>> next(iterator)
# 'triangle'
>>> next(iterator)
# 'square'
>>> next(iterator)
# 'pentagon'
>>> next(iterator)
# 'triangle'
>>> next(iterator)
# 'square'

Каждый раз, когда вызывается функция next(), возвращается следующее значение созданного итератора.


Итератор itertools.repeat() будет возвращать переданный объект (например словарь, список и т.д. ) снова и снова (бесконечно), если НЕ установить аргумент times, который отвечает за количество повторений.

Простой пример:

>>> from itertools import repeat
>>> lst = [1, 2]
# передаем объект списка
# и устанавливаем 2 повтора
>>> for list_obj in repeat(lst, 2):
>>>     for i in list_obj:
>>>         print(i)

# 1
# 2
# 1
# 2

Конечные итераторы.

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

>>> from itertools import accumulate
>>> list(accumulate(range(10)))
# [0, 1, 3, 6, 10, 15, 21, 28, 36, 45]

Здесь, в itertools.accumulate() передается диапазон из 10 чисел, 0-9. Функция accumulate() накапливает сумму по алгоритму: 0, 0+1, 1+2 и т. д.

Изменим поведение функции itertools.accumulate() по умолчанию и заставим накапливать значения путем умножения, для этого импортируем модуль operator.

>>> import itertools, operator
>>> list(itertools.accumulate(range(1, 5), operator.mul))
# [1, 2, 6, 24]

Таким образом, для каждой итерации - значения умножается (1x1, 1x2, 2x3 и т. д.), а не складывается.

В документации к itertools.accumulate() можно посмотреть другие интересные примеры, такие как амортизация ссуды или хаотичное рекуррентное соотношение.


Итератор itertools.chain() берет серию итераций и, по сути, "сглаживает" их в одну общую итерацию.

>>> from itertools import chain
>>> lst = ['foo', 'bar', ['one', 'two']]
>>> new_list = list(chain(lst, range(5))
>>> new_list
# ['foo', 'bar', ['one', 'two'], 0, 1, 2, 3, 4]

То же самое можно сделать без использования itertools при помощи оператора сложения или list.extend(), но при этом получится НЕ ИТЕРАТОР, а список, который увеличит расход памяти (будет содержать в себе копии списков):

>>> new_list = ['foo', 'bar', ['one', 'two']]
>>> new_list += range(5)
>>> new_list
# ['foo', 'bar', ['one', 'two'], 0, 1, 2, 3, 4]

# или
>>> new_list = ['foo', 'bar', ['one', 'two']]
>>> new_list.extend(range(5))
>>> new_list
# ['foo', 'bar', ['one', 'two'], 0, 1, 2, 3, 4]

Также можно использовать метод itertools.chain.from_iterable(). Он работает немного иначе, чем прямое использование itertools.chain() и принимает в качестве аргумента на серию итераций, а список с вложенными списками, которые необходимо сгладить.

Смотрим пример:

>>> from itertools import chain
>>> lst = ['foo', ['one', 'two', [1, 2]]]
>>> list(chain.from_iterable(lst))
# ['f', 'o', 'o', 'one', 'two', [1, 2]]

Итератор itertools.compress() полезен для фильтрации первой итерации на основе второй, которая должна быть списком логических значений.

Бессмысленный пример, зато показывает как работает функция:

>>> from itertools import compress
# есть строка
>>> letters = 'ABCDEFG'
# нужно получить итерацию из букв
>>> true_letter = 'AEG'
# создадим список логических значений
>>> bools = [x in true_letter for x in letters]
>>> bools
# [True, False, False, False, True, False, True]

# Подставляем
>>> list(compress(letters, bools))
# ['A', 'E', 'G']

Итератор itertools.dropwhile() отбрасывает элементы, пока критерий фильтра равен True. Другими словами, он не будет выдавать выходные данные, пока предикат не станет ложным. Это может увеличить время запуска, поэтому об этом нужно знать.

Смотрим на пример из документации:

>>> from itertools import dropwhile
>>> list(dropwhile(lambda x: x<5, [1,4,6,4,1]))
# [6, 4, 1]

Здесь в itertools.dropwhile() передается простая лямбда-функция, которая вернет значение True, если x меньше 5. Функция itertools.dropwhile() будет выполнять цикл по списку и передавать каждый элемент в лямбда-выражения. Если лямбда-выражение возвращает значение True, то это значение отбрасывается. Как только достигается число 6, лямбда-выражение возвратит значение False и сработает триггер, который разрешит вывод всех значений, следующих за этим числом.

Вместо лямбда-функции можно использовать обычную функцию, которая должна принимать значение из итерации и возвращать значение bool. Создадим функцию, которая возвращает True, если число больше 5.

>>> from itertools import dropwhile
>>> def trigger_to_five(x):
...     return x > 5 

>>> lst = [6, 7, 8, 9, 1, 2, 3, 10]
>>> list(dropwhile(trigger_to_five, lst))
# [1, 2, 3, 10]

Функция trigger_to_five() будет являться триггером. Она принимает все значения списка lst в порядке их следования и пока возвращаемые значения функции равны True, соответствующие значения последовательности будут удаляться из результирующей итерации.


Функция itertools.takewhile() в противоположна итератору itertools.dropwhile(), который рассматривали ранее. Функция itertools.takewhile() создаст итератор, который будет возвращать элементы из итерируемого объекта только до тех пор, пока предикат или триггер имеет значение True.

>>> from itertools import takewhile
>>> list(takewhile(lambda x: x<5, [1,4,6,4,1]))
# [1, 4]

Функция itertools.filterfalse() очень похожа на itertools.dropwhile(). Однако вместо триггера, который выводит все значения после его срабатывания, itertools.filterfalse() работает как фильтр и возвращает только те значения, которые оцениваются как False.

Воспользуемся функцией из предыдущего раздела:

>>> from itertools import filterfalse
>>> def greater_than_five(x):
...     return x > 5 
... 
>>> list(filterfalse(greater_than_five, [6, 7, 8, 9, 1, 2, 3, 10]))
# [1, 2, 3]

Итератор itertools.groupby() вернет последовательные ключи и группы из переданной последовательности. Это довольно сложно осмыслить, не видя примера.

from itertools import groupby

vehicles = [('Ford', 'Taurus'), ('Dodge', 'Durango'),
            ('Chevrolet', 'Cobalt'), ('Ford', 'F150'),
            ('Dodge', 'Charger'), ('Ford', 'GT')]

sorted_vehicles = sorted(vehicles, key=lambda make: make[0])

for key, group in groupby(sorted_vehicles, key=lambda make: make[0]):
    for make, model in group:
        print(f'{model} is made by {make}')
    print ("**** END OF GROUP ***\n")


# Cobalt is made by Chevrolet
# **** END OF GROUP ***
# Charger is made by Dodge
# Durango is made by Dodge
# **** END OF GROUP ***
# F150 is made by Ford
# GT is made by Ford
# Taurus is made by Ford
# **** END OF GROUP ***

Обратите внимание, что перед тем как передавать данные в функцию itertools.groupby(), они должны быть отсортированы по той же ключевой функции!


Инструмент itertools.starmap() создаст итератор, который может выполнять вычисления с использованием передаваемых функции и итерации. Разница между функциями map() и itertools.starmap() заключается в способе передачи аргументов вызываемой функции и аналогична разнице между function(a, b) и function(*c).

>>> from itertools import starmap
>>> def add(a, b):
...     return a+b
... 
>>> for item in starmap(add, [(2, 3), (4, 5)]):
...     print(item)

# 5
# 9

Функция itertools.tee(iter, n) создает n итераторов из одного итерируемого объекта iter. По умолчанию создается 2 итератора.

Смотрим как это работает:

>>> from itertools import tee
>>> data = 'ABC'
>>> iter1, iter2 = tee(data)
>>> for item in iter1:
...     print(item)

# A
# B
# C
>>> for item in iter2:
...     print(item)
... 
# A
# B
# C

Итератор itertools.zip_longest() можно использовать для объединения двух итерируемых объектов НЕ одинаково длины. Например, если одна итерация больше другой, то при объединении значений у более длинной последовательности отсутствующие элементы будут заполняться значением, переданным в аргумент fillvalue или None.

Смотрим пример:

>>> from itertools import zip_longest
>>> for item in zip_longest('ABCD', 'xy', fillvalue='BLANK'):
...     print (item)

# ('A', 'x')
# ('B', 'y')
# ('C', 'BLANK')
# ('D', 'BLANK')

Первые два кортежа представляют собой комбинации первых и вторых букв каждой строки соответственно. В последние два вставлено значение fillvalue.

Следует отметить, что если итерации, переданные в zip_longest, потенциально могут быть бесконечными, то необходимо обернуть функцию чем-то вроде itertools.islice(), чтобы ограничить количество вызовов.

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

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

Функция itertools.combinations() позволяет создать итератор со всеми возможными комбинациями элементов входной последовательности.

>>> from itertools import combinations
>>> list(combinations('WXYZ', 2))
[('W', 'X'), ('W', 'Y'), ('W', 'Z'), ('X', 'Y'), ('X', 'Z'), ('Y', 'Z')]

Чтобы сделать этот вывод более читабельным, объединим кортежи в одну строку.

>>> for item in combinations('WXYZ', 2):
...     print(''.join(item))

# WX
# WY
# WZ
# XY
# XZ
# YZ

Обратите внимание, что функция itertools.combinations() выполняет комбинации в лексикографическом порядке сортировки, следовательно, если передаваемая последовательность отсортирована, то и кортежи с комбинациями также будут отсортированы. Также стоит отметить, что , если все входные элементы уникальны, то и повторяющиеся значения в комбинациях создаваться не будут.


Итератор itertools.combinations_with_replacement() похож на itertools.combinations(). Единственное отличие заключается в том, что он фактически будет создавать комбинации, в которых элементы действительно повторяются.

Проиллюстрируем это предыдущим примером

>>> from itertools import combinations_with_replacement
>>> for item in combinations_with_replacement('WXYZ', 2):
...     print(''.join(item))

# WW
# WX
# WY
# WZ
# XX
# XY
# XZ
# YY
# YZ
# ZZ

Как можно видеть, теперь есть четыре новых элемента в выходных данных: WW, XX, YY и ZZ.


Функция itertools.product() предназначена для создания декартовых произведений из серии входных итераций.

>>> from itertools import product
>>> arrays = [(-1,1), (-3,3), (-5,5)]
>>> list(product(*arrays))
# [
#  (-1, -3, -5),
#  (-1, -3, 5),
#  (-1, 3, -5),
#  (-1, 3, 5),
#  (1, -3, -5),
#  (1, -3, 5),
#  (1, 3, -5),
#  (1, 3, 5)
# ]

Из кода можно заметить, что аргумент arrays распаковывается в момент передачи в функцию при помощи *. То есть, список будет "развернут" и последовательно применен к функции itertools.product(). Это означает, что передается три аргумента вместо одного.


Функция itertools.permutations() будет возвращать последовательные перестановки элементов указанной длины из итерации, которая передается в качестве аргумента. Подобно функции itertools.combinations(), перестановки выдаются в лексикографическом порядке сортировки.

>>> from itertools import permutations
>>> list(permutations('XYZ', 2))
# [
#    ('X', 'Y'), ('X', 'Z'), ('Y', 'X'), 
#    ('Y', 'Z'), ('Z', 'X'), ('Z', 'Y')
# ]
>>> for item in permutations('XYZ', 2):
...   print(''.join(item))
... 
# XY
# XZ
# YX
# YZ
# ZX
# ZY

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