9 Уровней применения функции zip в Python

Введение

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

Например, возьмём матрицу 2*3, представленную вложенным списком:

matrix = [[1, 2, 3], [1, 2, 3]]

И попробуем ответить на распространённый на собеседованиях по Python вопрос:

«Как транспонировать эту матрицу?».

Junior developer в ответ напишет несколько циклов for, в то время как senior ограничится всего одной строчкой кода:

matrix_T = [list(i) for i in zip(*matrix)]

Элегантно, согласитесь.

Если это однострочное решение пока что непонятно, не стоит переживать: дальше в статье мы подробно рассмотрим принцип действия мощной функции zip на девяти уровнях использования с полезными советами и хитрыми приёмами.

Если же такое решение вам знакомо, продолжайте читать статью и вы узнаете о ещё более крутых трюках и хитростях функции zip. 🍵

Уровень 0: как функция zip используется обычно

Функция zip объединяет элементы различных итерируемых объектов (таких как списки, кортежи или множества) и возвращает итератор.

Вот пример её применения для объединения двух списков:

id = [1, 2, 3, 4]
leaders = ['Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou']
record = zip(id, leaders)

print(record)
# <zip object at 0x7f266a707d80>

print(list(record))
# [(1, 'Elon Mask'), (2, 'Tim Cook'), (3, 'Bill Gates'), (4, 'Yang Zhou')]

Здесь функция zip возвращает итератор кортежей, где i-й кортеж содержит i-й элемент из каждого списка.

Принцип работы напоминает обычную застёжку-молнию.

Уровень 1: Zip работает с любым количеством итерируемых объектов

На самом деле функция zip в Python намного мощнее застёжки-молнии. Она имеет дело не с двумя, а с любым количеством итерируемых объектов одновременно.

Вот мы передаём в функцию zip один список:

id = [1, 2, 3, 4]
record = zip(id)
print(list(record))
# [(1,), (2,), (3,), (4,)]

А как насчёт трёх списков?

id = [1, 2, 3, 4]
leaders = ['Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou']
sex = ['male', 'male', 'male', 'male']
record = zip(id, leaders, sex)

print(list(record))
# [(1, 'Elon Mask', 'male'), (2, 'Tim Cook', 'male'), (3, 'Bill Gates', 'male'), (4, 'Yang Zhou', 'male')]

То есть неважно, сколько итерируемых объектов передаётся в функцию zip: она в любом случае работает как надо.

Кстати, если аргумента нет, функция zip возвращает пустой итератор.

Уровень 2: работа с неравными по длине аргументами

В реальности данные не всегда чистые и полные: иногда приходится иметь дело с неравными по длине итерируемыми объектами. По умолчанию результат функции zip берётся по длине самого короткого итерируемого объекта.

id = [1, 2]
leaders = ['Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou']
record = zip(id, leaders)

print(list(record))
# [(1, 'Elon Mask'), (2, 'Tim Cook')]

Так, в приведённом выше коде самый короткий список  —  это id. Поэтому record содержит только два кортежа, а два последних лидера в списке leaders были отброшены.

А что, если эти два последних лидера окажутся недовольны тем, как с ними поступили?

Python и здесь придёт на помощь. В модуле itertools есть функция zip_longest. Посмотрим, как с её помощью генерируется список record:

from itertools import zip_longest
id = [1, 2]
leaders = ['Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou']

long_record = zip_longest(id, leaders)
print(list(long_record))
# [(1, 'Elon Mask'), (2, 'Tim Cook'), (None, 'Bill Gates'), (None, 'Yang Zhou')]

long_record_2 = zip_longest(id, leaders, fillvalue='Top')
print(list(long_record_2))
# [(1, 'Elon Mask'), (2, 'Tim Cook'), ('Top', 'Bill Gates'), ('Top', 'Yang Zhou')]

Результат функции zip_longest основывается на самом длинном её аргументе. Дополнительный аргумент fillvalue, чьё значение по умолчанию равно None, помогает заполнить недостающие значения.

Уровень 3: операция распаковывания

Если в предыдущем примере получить сначала record, то как распаковать его на отдельные итерируемые объекты?

К сожалению, в Python нет функции распаковывания. Но если воспользоваться хитрыми приёмами звёздочек, распаковывание превращается в очень простую задачу.

record = [(1, 'Elon Mask'), (2, 'Tim Cook'), (3, 'Bill Gates'), (4, 'Yang Zhou')]
id, leaders = zip(*record)
print(id)
# (1, 2, 3, 4)
print(leaders)
# ('Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou')

С помощью звёздочки здесь выполнена операция распаковки: распакованы все четыре кортежа из списка record.

Если не прибегать к технике звёздочек, то же самое делается следующим методом:

record = [(1, 'Elon Mask'), (2, 'Tim Cook'), (3, 'Bill Gates'), (4, 'Yang Zhou')]

print(*record)  # распаковываем список одной звёздочкой
# (1, 'Elon Mask') (2, 'Tim Cook') (3, 'Bill Gates') (4, 'Yang Zhou')

id, leaders = zip((1, 'Elon Mask'), (2, 'Tim Cook'), (3, 'Bill Gates'), (4, 'Yang Zhou'))
print(id)
# (1, 2, 3, 4)
print(leaders)
# ('Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou')

Уровень 4: Создание и обновление словарей

С помощью функции zip очень просто создавать или обновлять dict, задействуя отдельные списки. Есть два однострочных решения:

  • использование dict comprehension и zip;
  • использование функции dict и zip.
id = [1, 2, 3, 4]
leaders = ['Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou']

# создаём словарь, используя «dict comprehension»
leader_dict = {i: name for i, name in zip(id, leaders)}
print(leader_dict)
# {1: 'Elon Mask', 2: 'Tim Cook', 3: 'Bill Gates', 4: 'Yang Zhou'}

# создаём словарь, используя функцию «dict»
leader_dict_2 = dict(zip(id, leaders))
print(leader_dict_2)
# {1: 'Elon Mask', 2: 'Tim Cook', 3: 'Bill Gates', 4: 'Yang Zhou'}

# обновляем
other_id = [5, 6]
other_leaders = ['Larry Page', 'Sergey Brin']
leader_dict.update(zip(other_id, other_leaders))
print(leader_dict)
# {1: 'Elon Mask', 2: 'Tim Cook', 3: 'Bill Gates', 4: 'Yang Zhou', 5: 'Larry Page', 6: 'Sergey Brin'}

В этом примере не используются циклы for. И насколько же элегантным он от этого становится, по-питоновски элегантным!

Уровень 5: функция zip вместо циклов for

Существует ещё один вариант применения, в котором функция zip заменяет цикл for,  —  это работа с последовательными элементами одной коллекции. Например, имеется список из целых чисел и надо вычислить разницу между соседними числами.

```

numbers = [12, 3, 7, 15, 8]

diff = [a-b for a, b in zip(numbers, numbers[1:])]

result: [9, -4, -8, 7]

Уровень 6: сортировка списков

Кроме того, с помощью zip выполняется сортировка двух связанных друг с другом списков:

```
>>> list1 = [3,2,4,1, 1]
>>> list2 = [‘three’, ‘two’, ‘four’, ‘one’, ‘one2’]
>>> list1, list2 = zip(*sorted(zip(list1, list2)))
>>> list1
(1, 1, 2, 3, 4)
>>> list2
(‘one’, ‘one2’, ‘two’, ‘three’, ‘four’)
```

Уровень 7: применение функции zip в циклах for

Типичный сценарий  —  когда имеешь дело с несколькими итерируемыми объектами одновременно. Функция zip показывает себя во всей красе, если задействовать её вместе с циклами for.

Продемонстрируем это на следующем примере:

products = ["cherry", "strawberry", "banana"]
price = [2.5, 3, 5]
cost = [1, 1.5, 2]
for prod, p, c in zip(products, price, cost):
    print(f'The profit of a box of {prod} is £{p-c}!')
# Прибыль от одного ящика вишни составляет 1,5 фунта!
# Прибыль от одного ящика клубники составляет 1,5 фунта!
# Прибыль от одного ящика бананов составляет 3 фунта!

Есть ли более элегантная реализация этого примера?

Уровень 8: транспонирование матрицы

Наконец, мы возвращаемся к распространённому на собеседованиях по Python вопросу:

«Как транспонировать матрицу?».

Теперь, когда мы уже имеем представление о функции zip, распаковывании с помощью звёздочки и list comprehensions, однострочное решение будет более чем очевидным:

matrix = [[1, 2, 3], [1, 2, 3]]
matrix_T = [list(i) for i in zip(*matrix)]
print(matrix_T)
# [[1, 1], [2, 2], [3, 3]]

Заключение

Функция zip в Python очень полезная и мощная. Правильное её использование помогает писать меньше кода, выполняя при этом больше операций. «Делай больше, используя меньше»  —  такова философия Python.

Читайте также:

Читайте нас в Telegram, VK и Яндекс.Дзен


Перевод статьи Yang Zhou: 7 Levels of Using the Zip Function in Python

Предыдущая статьяКак создать масштабируемую архитектуру для крупных мобильных проектов
Следующая статьяМутационное тестирование: создай мутанта и прокачай тест