Тур по по модулю itertools

В Python есть полезный модуль для работы с итерируемыми объектами под названием itertools. В этой статье мы расскажем о всех этих функциях, разделив их на шесть группы: бесконечные, комбинаторные, удлиняющие, фильтрующие, группирующие и вычисляющие.

Импортируем интересующий нас модуль:

>>> import itertools as it

Бесконечные функции из itertools

В модуле itertools есть три функции count, cycle и repeat, которые генерируют бесконечный итератор. Поэтому прерывать их нужно своими силами.

Функция count выдает числа, начиная со значения первого параметра. Вторым параметром можно указать шаг, который по умолчанию равен 1. Например, мы можем написать такой вот код на Python:

for num in it.count(-10, 5):
    print(num, end=" ")
"""
-5, -2.5, 0.0, 2.5, 5.0, ... и т.д.
"""

Функция cycle предназначена для итерирования около итерируемого объекта, переданного в качестве параметра:

for item in it.cycle("XYZ"):
    print(item, end=", ")
"""
X, Y, Z, X, Y, Z, ... и т.д.
"""

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

for i in it.repeat("Meow", 5):
    print(i, end="! ")

"""
Meow! Meow! Meow! Meow! Meow!
"""

Эти и последующие функции возвращают объект-итератор. Python не хранит их в виде списка. Вместо цикла for мы могли применять встроенную функцию next, которая бы возвращала сам элемент итератора.

Комбинаторные функции

Многие из нас со школы помнят задачи, которые требуют вычисления “це из эн по ка”. Они были связаны с комбинаторикой. Python предоставляет четыре комбинаторных функций:

• product размещение с повторениями, т.е. AA, BB учитываются. Считается по формуле n^k;
• permutations размещение без повторений, т.е. AA не учитываются. Считается по формуле n! / (n-k)!;
• combinations_wtih_replacement сочетания с повторениями, в этом случае AB != BA и BB учитываются. Считаются по формуле (n+k-1)! / (k! * (n-1)!).
• combinations сочетания без повторений, в этом случае AB == BA и BB не учитываются. Считается по формуле n! / (n-k)!

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

Удлиняющие функции

Функция chain объединяет несколько последовательностей в один. Например, у вас есть три списка, по каждому из них нужно пройтись и провести одну и ту же процедуру. Вы можете написать следующие код на Python:

l1 = [1, 2]; l2 = ["a", "b"]; l3 = [3, 4]
for item in it.chain(l1, l2, l3):
    print(item, end=", ")
"""
1, 2, a, b, 3, 4,
"""

Если же у вас вложенные списки, то используйте функцию chain.from_iterable.

l = ([1, 2], ["a", "b"], [3, 4])
for item in it.chain.from_iterable(l):
    print(item, end=", ")
"""
1, 2, a, b, 3, 4,

Функция tee возвращает заданное количество (по умолчанию 2) итераторов. Она может пригодится для тех случаев, когда может потребоваться именно итератор, а не весь список, причем понадобиться несколько раз.

l = [1, 2, 3]
for item in it.tee(l, 3):
    print(list(item))
"""
[1, 2, 3]
[1, 2, 3]
[1, 2, 3]
"""

Функция zip_longest может пригодится, когда один список больше по длине другого. Встроенная функция zip работает по самому короткому списку, т.е. цикл закончиться, когда исчерпается один из списков.
А вот функция zip_longest, наоборот, работает по самому длинному списку. Вместо значений исчерпанного списка Python будет подставлять значения None.

l1 = [1, 2, 3, 4]; l2 = ["a", "b"]
for el1, el2 in zip(l1, l2):
    print(el1, el2)
"""
1 a
2 b
"""

for el1, el2 in it.zip_longest(l1, l2):
    print(el1, el2)
"""
1 a
2 b
3 None
4 None
"""

Фильтрующие функции из itertools

Функции dropwhile будет брать в себя элементы, пока условие равно False. А функция takewhile будет брать в себя элементы пока условие будет равно True. Рассмотрим примеры на Python:

l = [4, 6, 4, 1, 5, 2]
is_even = lambda x: x % 2 == 0
for item in it.dropwhile(is_even, l):
    print(item, end=", ")
"""
1, 5, 2
"""

for item in it.takewhile(is_even, l):
    print(item, end=", ")
"""
4, 6, 4
"""

Значение 1, поданное функции дает ложь, поэтому начиная с него все остальные войдут в dropwhile или им закончатся в takewhile.

Функция filterfalse является противоположной встроенной функции filter. Если filter забирает те элементы, которые дают True, то filterfalse забирает те, которые дают False:

list(it.filterfalse(is_even, l))
"""
[1, 5]
"""
list(filter(is_even, l))
"""
[4, 6, 4, 2]
"""

Также есть функция compress, которая выбирает значения через селекторы. Селектором является последовательность, состоящая из значений True или False или им эквивалентам (например, все числа кроме нуля являются True).

list(it.compress("ABCDEF", [1,0,1,0,1,1]))
"""
['A', 'C', 'E', 'F']
"""

Можно объединить с другими функциями:

list(it.compress([1,2,3,4,5], it.cycle([1,0])))
"""
[1, 3, 5]
"""

Последняя функция из этого блока — islice. Она просто выбирает значения из заданного диапазона значения. У этой функции две сигнатуры: первая принимает два параметра (последовательность и сколько брать в себя (stop)), вторая принимает до четырех (начало, конец и шаг, по умолчанию равный 1). Причем ее можно использовать не только к итерируемым объектам, но и к итераторам. Примеры кода на Python:

# от 2 до 10 c шагом 2
list(it.islice('ABCDEFG', 2, 10, 2))
"""
['C', 'E', 'G']
"""

# Взять 5 элементов
list(it.islice(it.count(100, 15), 5))
"""
[100, 115, 130, 145, 160]
"""

Группирующие функции

В Python 3.10 появилась функция pairwise предназначена для создания пар из рядом стоящих элементов. Если количество элементов нечетное, то последний элемент встанет в пару с None.

Функция groupby создает группы из уникальных значений. Функция работает слева направо, если видит, что значение не совпадает с предыдущим, то заносит в себя. Причем она хранит, как сами элементы в виде итератора, так и их уникальные значения. Примеры:

keys = []
elements = []
for k, g in it.groupby("AAAABBBCCDAABBB"):
    keys.append(k)
    elements.append(list(g))

keys
"""
['A', 'B', 'C', 'D', 'A', 'B']
"""
elements
"""
[['A', 'A', 'A', 'A'], ['B', 'B', 'B'], ['C', 'C'], ['D'], ['A', 'A'], ['B', 'B', 'B']]
"""

Перед применением, возможно, следует сортировать последовательность.

Вычисляющие функции

Функция starmap применяет функцию к элементам внутри каждой вложенной последовательности:

list(it.starmap(pow, [(2,5), (3,2)]))
"""
[32, 9]
"""
list(it.starmap(max, [(1,2,3), (3,1,4)]))
"""
[3, 4]
"""

Функция accumulate является неким аналогом функции functools.reduce. Только accumulate сохраняет предыдущий результат:

list(it.accumulate([2, 3, 2], pow))
"""
[2, 8, 64]
"""

import functools as ft
ft.reduce(pow, [2, 3, 2])
"""
64
"""

 

Еще больше подробностей о применении этих функций для решения задач Data Science вы узнаете на наших образовательных курсах в лицензированном учебном центре обучения и повышения квалификации руководителей и ИТ-специалистов (менеджеров, архитекторов, инженеров, администраторов, Data Scientist’ов и аналитиков Big Data) в Москве:

• DPREP: Подготовка данных для Data Mining на Python
• Разработка и внедрение ML-решений
• Графовые алгоритмы. Бизнес-приложения
• Нейронные сети на Python
• NLP с Python

1. Документация к модулю