7 уровней использования функции zip в Python

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

В социальных сетях IT Resume мы уже рассматривали довольно много примеров того, как улучшить свой код на Python, как работать со строками в Python, обозревали основные функции библиотек и многое другое. А сегодня целую статью посвятим не всем известной, но очень полезной функции - zip.

В Python zip - встроенная функция, которая может сделать код очень элегантным. Однако ее использование не очень интуитивно понятно для новичков, а иногда и вызывает ошибки.

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

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

А вот и популярный вопрос на собеседовании по Python:

Как получить транспонирование указанной выше матрицы?

Junior разработчик может написать несколько циклов 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 в 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. Как следует из названия, это сестра функции zip, и ее результат зависит от самого длинного аргумента.

Воспользуемся 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')] 

Теперь ни Bill Gates, ни Yang Zhou не расстроятся. Необязательный аргумент fillvalue может помочь нам заполнить недостающие значения. Значение по умолчанию аргумента равно None, но мы можем поставить туда все, что угодно! Еще неизвестно, что было обиднее: не попасть в список record или получить рядом со своим именем что-нибудь эдакое :)

Уровень 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)  # unpack the list by one asterisk
 # (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

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

• Использование dict comprehension и zip

• Использование функции dict и zip

 id = [1, 2, 3, 4]
 leaders = ['Elon Mask', 'Tim Cook', 'Bill Gates', 'Yang Zhou']
 
 # create dict by 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'}
 
 # create dict by dict function
 leader_dict_2 = dict(zip(id, leaders))
 print(leader_dict_2)
 # {1: 'Elon Mask', 2: 'Tim Cook', 3: 'Bill Gates', 4: 'Yang Zhou'}
 
 # update
 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

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

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

 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}!')
 # The profit of a box of cherry is £1.5!
 # The profit of a box of strawberry is £1.5!
 # The profit of a box of banana is £3! 

Есть ли более аккуратный способ реализовать приведенный выше пример? Едва ли.

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

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

Как получить транспонированную матрицу?

Поскольку мы уже знакомы с функцией zip, распаковкой с помощью одной звездочки и dict comprehension, однострочное решение становится интуитивно понятным:

 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.