Разбираемся, как работает встроенная функция zip в Python, и пишем свою реализацию с помощью list comprehension

Расказывает Рювен Лёрнер, преподаватель


Многие слышали о функции zip в Python, а кто-то даже регулярно ей пользуется. Сегодня мы (из интереса и для общего развития) опишем, как можно реализовать её самому с помощью list comprehensions

Для начала поясню, что вообще делает функция zip, для тех, кто с ней раньше не сталкивался:

То есть функция берёт на вход несколько списков и создаёт из них список (в Python 3 создаётся не list, а специальный zip-объект) кортежей, такой, что первый элемент полученного списка содержит кортеж из первых элементов всех списков-аргументов. Таким образом, если ей передать три списка, то она отработает следующим образом:

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

Это весьма удобно, не находите? Каждый раз, когда я рассказываю о zip на своих уроках, у меня спрашивают о том, что будет, если в функцию передать массивы разной длины. Ответ простой — победит более короткий:

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

Есть одна возможность в Python, которая мне нравится даже больше, чем zip. Это списковое включение (англ. list comprehension). Именно поэтому, когда один из студентов недавно спросил меня, можем ли мы реализовать zip сами с помощью списковых включений, я просто не смог устоять.

Как же нам этого добиться? Начнём с первого, что приходит на ум:

В общем-то всё! Это работает. Но есть несколько моментов, которые всё же стоит доработать в этом методе.

Во-первых, оригинальная функция могла работать с массивами разной длины. Поэтому вместо range(len(s)) нам стоит использовать range(len(x)), где x — наиболее короткая последовательность. Для этого достаточно поместить все последовательности в один список, отсортировать этот список по длине элементов и выяснить длину элемента, оказавшегося под нулевым индексом:

Совмещаем это с предыдущим кодом:

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

Что осталось теперь? Как уже говорилось выше, Python 3 создаёт не список, а специальный zip-объект, возвращая итератор от него. Это сделано для того, чтобы код не ломался при обработке исключительно длинных последовательностей. Это можно реализовать, но уже не с помощью спискового включения (которое всегда возвращает список), а с помощью генератора. К счастью, для этого достаточно поменять квадратные скобки на круглые:

Готово! Мы реализовали свой полностью рабочий zip. Вы можете потренироваться и самостоятельно подумать, как ещё можно улучшить этот алгоритм.

Перевод публикации «Implementing “zip” with list comprehensions»