Python: enumerate() и zip()

Есть две функции, после знакомства с которыми часть циклов в Python начинает выглядеть заметно приятнее — enumerate() и zip().

Обе довольно простые, поэтому тут лучше сразу идти к примерам.

enumerate()

Допустим, есть список языков:

languages = ["Python", "Go", "Rust"]

Хотим вывести не только название, но и его номер.

Можно сделать так:

for i in range(len(languages)):
    print(i, languages[i])

Работает, но выглядит немного как обходной путь. Мы получаем длину списка, строим range(), затем используем индекс, чтобы снова обратиться к списку.

В Python для этого есть enumerate():

for i, language in enumerate(languages):
    print(i, language)

Результат:

0 Python
1 Go
2 Rust

По сути, enumerate() во время обхода добавляет к каждому элементу счётчик.

list(enumerate(languages))

[
    (0, "Python"),
    (1, "Go"),
    (2, "Rust"),
]

А запись:

for i, language in enumerate(languages):

просто распаковывает эти пары в две переменные.

Часто нумерация с нуля нам вообще не нужна. Например, если выводим рейтинг:

languages = ["Python", "Go", "Rust"]

for position, language in enumerate(languages, start=1):
    print(f"{position}. {language}")

Получится:

1. Python
2. Go
3. Rust

Мне кажется, это один из тех случаев, где код практически читается как обычный текст.

Кстати, число от enumerate() правильнее называть счётчиком, а не индексом. При обычном enumerate(languages) оно совпадает с индексом списка, но никто не мешает написать:

enumerate(languages, start=42)

Сам список от этого, конечно, не начнёт индексироваться с 42.

zip()

Теперь другая частая ситуация.

Есть имена:

names = ["Alice", "Bob", "Charlie"]

И отдельно возраст:

ages = [25, 31, 28]

Нужно пройти по ним одновременно.

Опять же, можно использовать индексы:

for i in range(len(names)):
    print(names[i], ages[i])

Но проще:

for name, age in zip(names, ages):
    print(name, age)

Получим:

Alice 25
Bob 31
Charlie 28

Название zip() здесь довольно удачное. Функция как будто застёгивает две последовательности вместе.

Alice    25
Bob      31
Charlie  28

Если посмотреть на результат через list():

list(zip(names, ages))

получим:

[
    ("Alice", 25),
    ("Bob", 31),
    ("Charlie", 28),
]

Причём последовательностей может быть сколько угодно:

names = ["Alice", "Bob", "Charlie"]
ages = [25, 31, 28]
cities = ["London", "Berlin", "Madrid"]

for name, age, city in zip(names, ages, cities):
    print(name, age, city)

Никакой работы с индексами здесь уже не требуется.

Есть один нюанс с zip()

Что произойдёт, если списки разной длины?

names = ["Alice", "Bob", "Charlie"]
ages = [25, 31]

for name, age in zip(names, ages):
    print(name, age)

Результат:

Alice 25
Bob 31

Charlie просто не попадёт в цикл.

zip() заканчивает работу вместе с самой короткой последовательностью.

Иногда именно это и нужно. А иногда это баг, который хотелось бы заметить.

Если длины обязаны совпадать, можно написать:

for name, age in zip(names, ages, strict=True):
    print(name, age)

Тогда при несовпадении длин получим ValueError.

Из zip() удобно делать словари

Есть список полей:

fields = ["name", "age", "city"]

И список значений:

values = ["Alice", 25, "London"]

Можно соединить их:

user = dict(zip(fields, values))

И получить:

{
    "name": "Alice",
    "age": 25,
    "city": "London",
}

Простой приём, но периодически оказывается очень кстати.

А можно использовать их вместе

Допустим, у нас есть участники и их результаты:

names = ["Alice", "Bob", "Charlie"]
scores = [95, 87, 91]

Хотим вывести:

1. Alice — 95
2. Bob — 87
3. Charlie — 91

Тогда можно совместить обе функции:

for position, (name, score) in enumerate(
    zip(names, scores),
    start=1,
):
    print(f"{position}. {name} — {score}")

На первый взгляд эта конструкция может выглядеть немного странно:

position, (name, score)

Но если посмотреть на данные, всё становится понятнее.

zip() даёт:

("Alice", 95)
("Bob", 87)
("Charlie", 91)

А enumerate() добавляет к каждой такой паре номер:

(1, ("Alice", 95))
(2, ("Bob", 87))
(3, ("Charlie", 91))

Отсюда и такая распаковка:

position, (name, score)

Главная идея здесь та же, что и во многих других возможностях Python: если приходится вручную управлять индексами, стоит сначала проверить, нет ли способа выразить задачу напрямую.