Автор: Во всем, что мы делаем, есть некоторые итерационные проблемы (и их решения). В стандартной библиотеке Python Itertools мы можем найти набор многих очень полезных функций, создающих итераторы. Все из соображений, чтобы не изобретать велосипед и просто создавать код быстрее 🙂 .
Содержание
Встроенные итераторы в Python Itertools
Стандартная библиотека Python известна тем, что она огромна. Мы действительно можем делать многие вещи без привлечения внешних зависимостей. Это, несомненно, большой плюс. В этом посте я познакомлю вас с тем, что такое Python Itertools. То есть, набор из 19 функций, которые создают интересные итераторы.
Разделение итераторов в библиотеке itertools
В официальной документации вы найдете описание каждого итератора в отдельности, а также его наглядную реализацию.
Кроме того, итераторы делятся на три категории: бесконечные итераторы, итераторы, которые заканчиваются на самой короткой входной последовательности, и комбинаторные итераторы.
Давайте рассмотрим их подробнее.
Бесконечные итераторы
1. COUNT
Создает итератор, который возвращает значения, начиная с того, которое задано в качестве начального параметра. Последующие увеличиваются на величину шага – второй параметр.
Пример использования:
count_iterator = count()
for i in range(3):
print(next(count_iterator))
# 0
# 1
# 2Code language: PHP (php)Мы можем изменить начальное значение:
count_iterator = count(1)
for i in range(3):
print(next(count_iterator))
# 1
# 2
# 3Code language: PHP (php)Также мы можем изменить величину, на которую изменяются последовательные числа. В приведенном выше случае мы возвращаем только четные числа:
count_iterator = count(step=2)
for i in range(3):
print(next(count_iterator))
# 0
# 2
# 4Code language: PHP (php)Конечно, мы можем установить отрицательный шаг:
count_iterator = count(1, step=-1)
for i in range(3):
print(next(count_iterator))
# 1
# 0
# -1Code language: PHP (php)Используя count, мы также можем повторить работу встроенной функции enumerate:
colors = ['red', 'blue', 'green']
indexed_colors = [color for color in zip(count(), colors)]
print(indexed_colors)
# [(0, 'red'), (1, 'blue'), (2, 'green')]Code language: PHP (php)2. CYCLE
cycle принимает в качестве параметра итерабельную переменную. Этот итератор проходит через переданный аргумент iterable. Разница в том, что если он закончится (дойдет до конца), он все равно вернет значения из начала. Это происходит потому, что cycle хранит каждый элемент в памяти.
cycle_iterator = cycle('кассир')
for i in range(8):
print(next(cycle_iterator))
# к
# а
# с
# с
# и
# р
Code language: PHP (php)Другой пример – разделить участников на X групп, используя правило “последовательный отсчет до X”.
users = ['Иван', 'Петр', 'Евгений', 'Александр', 'Михаил', 'Алексей', 'Андрей', 'Сергей']
cycle_iterator = cycle([1, 2, 3, 4])
users_teams = [name for name in zip(cycle_iterator, users)]
print(users_teams)
# [(1, 'Иван'), (2, 'Петр'), (3, 'Евгений'), (4, 'Александр'), (1, 'Михаил'), (2, 'Алексей'), (3, 'Андрей'), (4, 'Сергей')]Code language: PHP (php)3. REPEAT
Эта функция возвращает заданный объект X раз. Если параметр times не передан, объект будет возвращен бесконечно длинным.
В официальной документации объясняется, что repeat в основном используется как постоянный аргумент, передаваемый функции map. Его также можно использовать для добавления постоянного значения к кортежу.
Следующий пример взят из официальной документации и вычисляет квадрат силы для чисел от 0 до 9.
squares = list(map(pow, range(10), repeat(2)))
print(squares) # [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]Code language: PHP (php)Итераторы, заканчивающиеся кратчайшей входной последовательностью
4. ACCUMULATE
Это функция, которая создает накопленные суммы или результаты двоичных функций. Она ведет себя аналогично известной функции reduce из библиотеки functools, за исключением того, что возвращает текущие значения для каждого элемента, переданного в функцию.
Простой пример сложения последовательных итерируемых элементов вместе:
numbers = [1, 2, 3, 4, 5]
result = list(accumulate(numbers))
print(result) # [1, 3, 6, 10, 15]Code language: PHP (php)По умолчанию значения складываются друг с другом. Если вы хотите изменить операцию, вы можете сделать это, передав функцию в качестве параметра func.
Помните, что она должна принимать ровно два аргумента (так называемая бинарная функция). Например, пусть это будет продукт:
numbers = [1, 2, 3, 4, 5]
result = list(accumulate(numbers, func=operator.mul))
print(result) # [1, 2, 6, 24, 120]Code language: PHP (php)Другим примером может быть возврат текущего максимального значения:
numbers = [1, 2, 1, 4, 3]
result = list(accumulate(numbers, func=max))
print(result) # [1, 2, 2, 4, 4]Code language: PHP (php)Количество элементов в возвращаемом итераторе будет равно количеству элементов в итераторе, переданном в качестве аргумента.
Единственным исключением является передача необязательного аргумента initial, который задает начальное значение.
numbers = [1, 2, 1, 4, 3]
result = list(accumulate(numbers, func=max, initial=3))
print(result) # [3, 3, 3, 3, 4, 4]Code language: PHP (php)5. CHAIN
Очень полезная функциональность для итерации по нескольким наборам один за другим. Сначала выполняется итерация над первой итерируемой, затем переходим ко второй и так далее.
colors = ('red', 'orange')
numbers = (1, 2)
for element in chain(colors, numbers):
print(element)
# red
# orange - конец первой итерации
# 1 - начало второй итерации
# 2Code language: PHP (php)6. CHAIN.FROM_ITERABLE
Описанная выше функция chain имеет метод from_iterable, наиболее популярное применение которого, возможно, заключается в уплощении множества.
input_list = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
result = chain.from_iterable(input_list)
print(list(result)) # [1, 2, 3, 4, 5, 6, 7, 8, 9]Code language: PHP (php)7. COMPRESS
Принимает два итерируемых набора в качестве параметров. Итератор, возвращаемый compress, возвращает только те значения из первого набора, для которых значение во втором наборе ( selectors) на той же позиции является логической истиной (True).
data = ['h', 'b', 'e', 'o', 'b', 'p']
selectors = [0, 1, 0, 1, 1, 0]
result = compress(data, selectors)
print(''.join(list(result))) # bobCode language: PHP (php)Если набор данных больше, чем набор селекторов, избыточные элементы автоматически имеют значения False – они не появятся в результирующем итераторе.
8. DROPWHILE
Возвращает итератор, который игнорирует/отбрасывает все значения из итерируемого набора до тех пор, пока первый элемент не выполнит условие (параметр predicate).
С этого момента каждый последующий элемент будет возвращен, даже если он не удовлетворяет условию, поскольку он больше не проверяется.
input = [1, 2, 4, 6, 2, 1]
result = dropwhile(lambda x: x < 3, input)
print(list(result)) # [4, 6, 2, 1]Code language: PHP (php)9. TAKEWHILE
Он ведет себя прямо противоположно описанному выше dropwhile. Он возвращает значения до тех пор, пока выполняется условие. С момента первого нарушения условия (параметр predicate) итератор не будет возвращать другое значение.
input = [1, 2, 4, 6, 2, 1]
result = takewhile(lambda x: x < 3, input)
print(list(result)) # [1, 2]Code language: PHP (php)10. FILTERFALSE
Создает итератор, который возвращает только те значения, которые не удовлетворяют условию из параметра предиката. Если аргумент предиката равен None, то возвращаются только значения, которые переводятся в boolean false.
Пример фильтрации отрицательных цифр:
input = [1, -1, -2, 4, 5, -6]
result = filterfalse(lambda x: x < 0, input)
print(list(result)) # [1, 4, 5]Code language: PHP (php)11. GROUPBY
Как следует из названия, эта функция группирует последовательные элементы из переданного множества. Группы, которые он возвращает, являются итераторами, поэтому, если мы хотим использовать их более одного раза, стоит превратить их в список.
input = 'AAAABBBCCDYYY'
for key, group in groupby(input):
print(f'{key} - {list(group)}')
# A - ['A', 'A', 'A', 'A']
# B - ['B', 'B', 'B']
# C - ['C', 'C']
# D - ['D']
# Y - ['Y', 'Y', 'Y']Code language: PHP (php)Очень важно, чтобы входные данные были отсортированы. Почему? Потому что groupby создает новую группу, если следующее значение, потребляемое ею, отличается от предыдущего. При передаче несортированных данных результат будет выглядеть следующим образом:
input = 'ABBABBAA'
for key, group in groupby(input):
print(f'{key} - {list(group)}')
# A - ['A']
# B - ['B', 'B']
# A - ['A']
# B - ['B', 'B']
# A - ['A', 'A']Code language: PHP (php)То есть мы получаем дублирующие группы. В отличие от этого, GROUP BY, известный из SQL, может группировать данные, даже если они не отсортированы, в то время как groupby требует отсортированных данных на входе.
groupby принимает аргумент key для указания значения, которое будет использоваться для группировки данных.
input = [
{'name': 'Иван', 'level': 1},
{'name': 'Андрей', 'level': 3},
{'name': 'Сергей', 'level': 4},
{'name': 'Михаил', 'level': 4},
]
for key, group in groupby(input, key=lambda x: x['level']):
print(f'{key} - {list(group)}')
# 1 - [{'name': 'Иван', 'level': 1}]
# 3 - [{'name': 'Андрей', 'level': 3}]
# 4 - [{'name': 'Сергей', 'level': 4}, {'name': 'Михаил', 'level': 4}Code language: PHP (php)12. ISLICE
Он ведет себя аналогично разбиению, например, списков (islice возвращает итератор). Мы передаем iterable с опциями: начальный индекс, конечный индекс и шаг.
Стоит помнить, что если мы передаем только один аргумент, кроме самого итератора, то он будет обозначать конечную позицию возвращаемого итератора.
result = islice(range(5), 3)
print(list(result)) # [0, 1, 2]Code language: PHP (php)Добавляя еще один аргумент, мы по очереди (в таком порядке) указываем начальную и конечную позиции.
result = islice(range(5), 2, 4)
print(list(result)) # [2, 3]Code language: PHP (php)Кроме того, мы можем указать шаг/переход, который означает, какое значение будет возвращено.
result = islice(range(10), 1, 9, 3)
print(list(result)) # [1, 4, 7]Code language: PHP (php)13. STARMAP
Она работает очень похоже на известную функцию map с одним небольшим отличием – в качестве второго параметра она принимает итерируемое множество.
input = [(1, 0, -1), (4, 2, 3), (5, 2, 7)]
result = starmap(max, input)
print(list(result)) # [1, 4, 7]Code language: PHP (php)То есть, вместо того, чтобы каждый раз выполнять функцию map для следующего кортежа, мы передаем набор кортежей. starmap выполняет итерации над каждым кортежем и выполняет над ним функцию – в нашем случае функцию, которая выбирает наибольшее значение.
14. TEE
Как мы знаем вы можете пройти через созданный итератор только один раз. После исчерпания вы должны создать еще один, если хотите использовать его снова.
Функция tee позволяет нам создать X независимых итераторов из одной итерации. По умолчанию он создает два итератора.
text = 'abc'
iterator1, iterator2 = tee(text)
print(next(iterator1)) # a
print(next(iterator1)) # b
print(next(iterator1)) # c
print(next(iterator2)) # a
print(next(iterator2)) # b
print(next(iterator2)) # cCode language: PHP (php)Важно: в документации сказано, что после использования функции tee объект iterable больше нигде не должен использоваться. Это может привести к неожиданному и ошибочному поведению.
Кроме того, в официальной документации содержится информация, предупреждающая, что tee() не является потокобезопасным. Одновременное использование итераторов, возвращаемых tee, может привести к исключению RuntimeError.
15. ZIP_LONGEST
Если вы использовали zip для объединения двух итераторов, то вы можете знать, что количество элементов такого итератора равно кратчайшему итератору. Доказательство:
input_1 = [1, 2, 3, 4]
input_2 = 'abc'
result = zip(input_1, input_2)
print(list(result)) # [(1, 'a'), (2, 'b'), (3, 'c')]Code language: PHP (php)Как видите, итератор вернул только три элемента, потому что именно столько элементов имеет самый короткий набор (input_2). zip_longest ведет себя противоположным образом. Возвращает количество элементов, равное самой длинной из переданных итераций.
По умолчанию отсутствующие значения из более короткого набора будут заменены на None. Мы можем передать любое другое значение, добавив аргумент fillvalue.
input_1 = [1, 2, 3, 4]
input_2 = 'abc'
result = zip_longest(input_1, input_2, fillvalue=0)
print(list(result)) # [(1, 'a'), (2, 'b'), (3, 'c'), (4, 0)]Code language: PHP (php)Комбинаторные итераторы
16. PRODUCT
Он используется для вычисления декартова произведения. Если входные итералы отсортированы, то выходные элементы также будут отсортированы.
Ниже я приведу очень интересный пример. Достаточно простым способом мы генерируем все игральные карты.
FACE_CARDS = ("J", "Q", "K", "A")
SUITS = ("♥", "♦", "♣", "♠")
DECK = list(
product(
chain(range(2, 11), FACE_CARDS),
SUITS,
)
)
for card in DECK:
print("{:>2}{}".format(*card), end=" ")
if card[1] == SUITS[-1]: # переход на новую строку, если пики
print()
# 2♥ 2♦ 2♣ 2♠
# 3♥ 3♦ 3♣ 3♠
# 4♥ 4♦ 4♣ 4♠
# 5♥ 5♦ 5♣ 5♠
# 6♥ 6♦ 6♣ 6♠
# 7♥ 7♦ 7♣ 7♠
# 8♥ 8♦ 8♣ 8♠
# 9♥ 9♦ 9♣ 9♠
# 10♥ 10♦ 10♣ 10♠
# J♥ J♦ J♣ J♠
# Q♥ Q♦ Q♣ Q♠
# K♥ K♦ K♣ K♠
# A♥ A♦ A♣ A♠Code language: PHP (php)Кроме того, в приведенном выше примере мы использовали ранее описанную цепочку, чтобы легко сложить пронумерованные карты с валетом, королевой, королем и тузом. Благодаря product() мы получили комбинации всех фигур с цветами.
Кроме того, product позволяет указать, сколько раз дублировать набор, переданный в качестве аргумента. Это делается с помощью параметра repeat. Вот простой пример, демонстрирующий поведение этого параметра:
result = product(range(2))
print(list(result)) # [(0,), (1,)]Code language: PHP (php)Как вы видите, мы передали только одну итерабельную переменную, поэтому на самом деле мы получим те же значения, которые передали. Теперь добавим параметр repeat=2:
result = product(range(2), repeat=2)
print(list(result)) # [(0, 0), (0, 1), (1, 0), (1, 1)]Code language: PHP (php)Сразу стало лучше. Если мы захотим вычислить декартово произведение, используя те же наборы, полезно запомнить этот параметр.
17. PERMUTATIONS
Благодаря этой функции мы получаем перестановки переданного множества. По умолчанию длина каждой перестановки равна длине переданной итерабельной переменной. Передавая аргумент r, мы можем манипулировать им.
result = permutations('ABC')
print(list(result)) # [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]
result = permutations('ABC', r=2)
print(list(result)) # [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]Code language: PHP (php)18. COMBINATIONS
Вычисляет комбинации для заданного набора. В этом случае необходимо передать аргумент r, указывающий длину комбинации.
result = combinations('ABC', r=2)
print(list(result)) # [('A', 'B'), ('A', 'C'), ('B', 'C')]Code language: PHP (php)19. COMBINATIONS_WITH_REPLACEMENT
Он ведет себя аналогично combinations, описанным выше, за исключением того, что позволяет повторять один и тот же элемент.
result = combinations_with_replacement('ABC', r=2)
print(list(result)) # [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]Code language: PHP (php)Не только Python Itertools…
Если вы еще не нашли то, что искали в стандартной библиотеке Python Itertools, вы можете найти это в пакете more-itertools. Он содержит множество различных и, главное, эффективных инструментов, которые могут сэкономить нам много времени.
Стоит просмотреть список функций, потому что вы можете найти что-то для себя. На момент написания этой заметки проект активно развивается.
