dict comprehension: словарь за строку

Словари в этом разделе собирались по кусочку: пара за парой в литерале, добавлением через присваивание, пачкой через update и |. Но чаще всего словарь строится из данных, которые уже есть: из списка чисел, из старого словаря, из пар, перебранных в цикле. Для этого случая у Python есть короткая запись — dict comprehension: цикл, свёрнутый в одну строку, который собирает словарь целиком.

От цикла к comprehension: квадраты чисел

squares = {n: n * n for n in range(1, 7)}

print(squares)

{1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36}

Читается запись по частям. Внутри фигурных скобок стоит пара n: n * n — что положить в словарь: слева от двоеточия ключ, справа значение. Дальше for n in range(1, 7) — откуда брать элементы: comprehension пройдёт по каждому n и для каждого вычислит пару. Тот же словарь циклом выглядел бы так: пустой словарь, for, присваивание squares[n] = n * n — пять строк против одной, и разница растёт с каждой задачей.

Цены со скидкой: преобразование и фильтр

prices = {'молоко': 89, 'хлеб': 42, 'сыр': 590}

sale = {name: round(price * 0.9) for name, price in prices.items()}
print(sale)

{'молоко': 80, 'хлеб': 38, 'сыр': 531}

prices = {'молоко': 89, 'хлеб': 42, 'сыр': 590, 'чай': 310}

cheap = {name: price for name, price in prices.items() if price < 100}
print(cheap)

{'молоко': 89, 'хлеб': 42}

В обоих случаях после for стоят два имени: перебор идёт по items, поэтому каждая итерация распаковывается в name и price. До for записывают, что положить в пару, после for — из чего берутся элементы, а условие после if отсеивает лишнее. Скидка меняет значение — выражение слева может быть любым вычислением; фильтр отбирает пары целиком — условие возвращает False, и пара в новый словарь не попадает. Исходный prices при этом не тронут: comprehension всегда строит новый словарь.

Инверсия: ключ и значение меняются местами

phone = {'Анна': '101', 'Борис': '102', 'Вера': '103'}

by_number = {number: name for name, number in phone.items()}
print(by_number)
print(by_number['103'])

{'101': 'Анна', '102': 'Борис', '103': 'Вера'}
Вера

Инверсия, которая не теряет данные

Раз замена значений ключами может съесть данные, нужен обратный ход: собрать словарь, где по значению лежит не одно имя, а список всех имён с таким значением. Здесь ключ остаётся уникальным — это само значение, — а потерять несколько имён уже невозможно. Синтаксис получается плотнее, но смысл тот же: один проход вместо цикла с проверками.

salaries = {'Анна': 90, 'Борис': 90, 'Вера': 120}

by_salary = {}
for name, value in salaries.items():
    by_salary.setdefault(value, []).append(name)

print(by_salary)

{90: ['Анна', 'Борис'], 120: ['Вера']}

Разбор по шагам. by_salary.setdefault(value, []) достаёт список по ключу value, а если его нет — создаёт пустой и кладёт на место. Дальше .append(name) добавляет имя в этот список. Получается группировка: по ключу-зарплате собирается список имён. Тот же результат даёт defaultdict(list) из урока 12 — короче, потому что не нужно проверять наличие ключа.

from collections import defaultdict

salaries = {'Анна': 90, 'Борис': 90, 'Вера': 120}

grouped = defaultdict(list)
for name, value in salaries.items():
    grouped[value].append(name)

print(dict(grouped))

{90: ['Анна', 'Борис'], 120: ['Вера']}

Одно отличие, которое стоит знать: defaultdict создаёт пустой список при обращении к отсутствующему ключу, даже если вы просто читаете. Поэтому в конце мы печатаем dict(grouped) — обычный словарь, без побочного эффекта. Если бы ключ '130' не встречался в данных, но мы его запросили, в grouped появился бы лишний пустой список, и длина словаря разъехалась бы с числом реальных групп.

Существует и однострочная версия той же группировки — вложенный comprehension. Работает она медленнее цикла (внутренний список перебирает весь словарь для каждого значения), но записывается так:

salaries = {'Анна': 90, 'Борис': 90, 'Вера': 120}

grouped = {value: [name for name, v in salaries.items() if v == value]
           for value in salaries.values()}

print(grouped)

{90: ['Анна', 'Борис'], 120: ['Вера']}

Читается в два слоя. Внешний comprehension перебирает salaries.values() — то есть уникальные зарплаты, — и для каждой строит список имён. Внутренний comprehension отбирает все пары, у которых зарплата совпадает с текущей. Результат тот же, что у цикла, и данные не теряются. Но если в словаре 50 тысяч пар, внутренний перебор будет выполнен 50 тысяч раз — вот здесь цикл с setdefault уже выигрывает, и разница станет заметной. Правило простое: пара-три тысячи записей — comprehension, больше — цикл.

Словарь из списка: enumerate и zip

Comprehension умеет брать не только из словаря. Если под рукой список, ключ можно построить из индекса — и вот тут пригодится enumerate: он отдаёт пары «номер, элемент», нумеруя с нуля.

words = ['сок', 'компот', 'чай', 'кофе']

print({i: w for i, w in enumerate(words)})
print({w: i for i, w in enumerate(words)})

{0: 'сок', 1: 'компот', 2: 'чай', 3: 'кофе'}
{'сок': 0, 'компот': 1, 'чай': 2, 'кофе': 3}

Второй вариант — инверсия списка: по слову находится его номер. Обратите внимание на порядок переменных в comprehension: for i, w in enumerate(words) даёт сначала индекс, потом элемент, и в фигурных скобках вы вправе поставить их как угодно. В первой строке индекс стал ключом, во второй — значением. Если нужен нумерация с единицы, добавьте start=1 — часто это удобнее для нумерации строк таблицы.

print({i: letter for i, letter in enumerate('абв', start=1)})

{1: 'а', 2: 'б', 3: 'в'}

Второй случай — два списка одинаковой длины, которые надо разложить по парам. Здесь работает zip: он соединяет списки по позициям, и получается готовый словарь «имя — зарплата».

names = ['Анна', 'Борис']
salaries = [90, 120]

print({n: s for n, s in zip(names, salaries)})

{'Анна': 90, 'Борис': 120}

Оговорка про zip, на которой спотыкаются: он останавливается на самом коротком списке. Если в одном из списков окажется лишний элемент, он молча потеряется — никакой ошибки, никакого предупреждения. Проверяйте длины заранее, если данные приходят из разных источников: assert len(names) == len(salaries). Для словарей вместо списков zip работает ещё удобнее: он соединяет пары items() напрямую.

Ключ из значения: когда имена у списка нет

Иногда словарь нужен не для красоты, а чтобы обратиться по другому полю. Пример из практики: в логе лежит список пар «код — имя», а искать надо по имени. Comprehension разворачивает список пар в словарь, где ключ — код, а значение — имя; перевернуть можно вторым проходом.

pairs = [(1, 'Анна'), (2, 'Борис')]

print({name: code for code, name in pairs})

{'Анна': 1, 'Борис': 2}

Здесь for code, name in pairs распаковывает кортежи прямо в цикле, и слева от двоеточия стоит name — то есть второй элемент кортежа. В исходном списке первым шёл код, а стал он значением: имя теперь ключ, по которому ищем код. Обратный вариант получится простой перестановкой в фигурных скобках: {code: name for code, name in pairs} даст словарь «код — имя». Проверить, что разворот прошёл правильно, можно сравнением длин: у списка пар длина равна числу кортежей, а у словаря — числу уникальных ключей, и если в списке были повторы имён, словарь получится короче.

Тернарный оператор: значение зависит от условия

До сих пор условие в comprehension только отсеивало пары: не прошёл — не попал в словарь. А если значение должно зависеть от условия, нужен тернарный оператор — компактный если то, иначе то. Пишется он прямо в выражении значения.

prices = {'молоко': 89, 'хлеб': 42, 'сыр': 590}

sale = {name: round(price * 0.9) if price > 50 else price
        for name, price in prices.items()}

print(sale)

{'молоко': 80, 'хлеб': 42, 'сыр': 531}

Читается так: если цена больше 50, в пару кладётся округлённая цена со скидкой, иначе — сама цена без изменений. Хлеб за 42 остался нетронутым, молоко со 89 превратилось в 80, сыр — в 531. Тернарный оператор — это выражение, его значение можно использовать дальше как обычное число. Второй вариант записи, круглыми скобками: {name: (round(price * 0.9) if price > 50 else price) for ...} — так его ещё проще отличить от условия-фильтра после for.

Когда comprehension — уже слишком

Comprehension — не лучший стиль на все случаи. Как только внутри появляются два уровня вложенности, несколько условий или вычисление с приоритетом, строка перестаёт читаться с одного взгляда. Правило простое: если запись помещается в одну строку шириной экрана — оставляйте, если нет — возвращайтесь к обычному циклу. Читаемость важнее компактности.

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

cells = {(r, c): matrix[r][c]
         for r in range(len(matrix))
         for c in range(len(matrix[r]))}

print(cells)

{(0, 0): 1, (0, 1): 2, (1, 0): 3, (1, 1): 4}

Здесь получился словарь-таблица: ключ-кортеж (строка, столбец), значение — ячейка. Такой приём удобен для разреженных матриц и для игровых полей. Запись разбита на три строки, и каждый кусок на своей — так её ещё можно читать. Но если бы в comprehension понадобился ещё один уровень, мы бы честно переписали это обычным циклом: код, который нужно читать дважды, должен читаться легко с первого раза.

Что дальше

Comprehension собирает словарь из данных, которые уже в программе. Дальше — словарь, который приходит снаружи: в уроке 16 разбираем JSON и json.loads, который превращает текст в словарь, а json.dumps — обратно. А если данные надо сначала разобрать на части, посмотрите урок про вложенные словари — там цепочка ключей и безопасное чтение через get.