Пересечение нескольких множеств в одну строчку

Автор: Илья Щуров
Категория: короткие заметки

Теги: python, set, intesection, one-liner

Сегодня на паре по программированию для журналистики данных мы столкнулись с такой задачей: есть список множеств, нужно найти пересечение их всех. Как это проще всего сделать в Python? Оказывается, у этой задачи есть решение в одну строчку.

Собственно, дело было так. Мы смотрели на данные по госзакупкам через API сайта clearspending.ru. Там у каждого контракта куча разных атрибутов. Как видно из следующих примеров, наборы атрибутов (в данном случае, ключей в словаре) у разных контрактов различаются. Мы хотели найти такие атрибуты, которые есть у всех контрактов.

In [2]:
import requests
# пример из документации
# https://github.com/idalab/clearspending-examples/wiki/Описание-API-Контракты
entrypoint = "http://openapi.clearspending.ru/restapi/v3/contracts/search/"
r = requests.get(entrypoint, {'customerregion': '05', 'sort':'-price'})
response = r.json()
contracts = response['contracts']['data']
In [3]:
contracts[0].keys()
Out[3]:
dict_keys(['finances', 'documentBase', 'versionNumber', 'mongo_id', 'suppliers', 'placingWayCode', 'contractUrl', 'foundation', 'products', 'scan', 'fileVersion', 'contractProcedure', 'economic_sectors', 'signDate', 'fz', 'currentContractStage', 'printFormUrl', 'price', 'protocolDate', 'number', 'regNum', 'currentContractStage_raw', 'loadId', 'attachments', 'customer', 'placing', 'regionCode', 'id', 'currency', 'singleCustomerReason', 'execution', 'publishDate', 'schemaVersion'])
In [4]:
contracts[1].keys()
Out[4]:
dict_keys(['signDate', 'versionNumber', 'mongo_id', 'finances', 'contractUrl', 'foundation', 'misuses', 'regionCode', 'fileVersion', 'documentBase', 'fz', 'currentContractStage', 'price', 'placing', 'number', 'products', 'publishDate', 'customer', 'regNum', 'suppliers', 'id', 'currency', 'execution', 'loadId'])

По такому поводу я рассказал про множества в Python (давно откладывал это — к слову не приходилось) и написал такой код.

In [5]:
fields = set(contracts[0])
# сделали множество из списка ключей первого контракта
# здесь словарь contracts[0] рассматривается как iterable
# а в этом случае он итерирует свои ключи
# поэтому .keys() дописывать не нужно
for contract in contracts[1:]:
    fields.intersection_update(contract)
    # пересечь множество fields с множеством, полученным из списка ключей
    # очередного контракта и результат записать в fields
    # иными словами, выкинуть из fields те элементы, которых нет в
    # списке ключей очередного контракта
fields
Out[5]:
{'currency',
 'customer',
 'fileVersion',
 'fz',
 'id',
 'loadId',
 'mongo_id',
 'number',
 'price',
 'products',
 'publishDate',
 'regNum',
 'regionCode',
 'versionNumber'}

Он мне перестал нравиться ещё до того, как я закончил его писать. Ну, право дело, ведь когда нам нужно сложить числа в списке, мы не пишем цикл — мы просто вызываем функцию sum(). Должно, наверное, и для множеств быть что-то похожее? На паре тратить время на поиски не хотелось, но, придя домой, я всё же решил найти ответ на этот вопрос. Оказывается, есть очень просто решение!

In [6]:
fields = set.intersection(*[set(contract) for contract in contracts])
fields
Out[6]:
{'currency',
 'customer',
 'fileVersion',
 'fz',
 'id',
 'loadId',
 'mongo_id',
 'number',
 'price',
 'products',
 'publishDate',
 'regNum',
 'regionCode',
 'versionNumber'}

Дело в том, что set.intersection() принимает на вход любое количество аргументов! С помощью спискового включения мы делаем список множеств, составленных из ключей каждого контракта, затем звёздочкой «распаковываем» этот список в набор аргументов set.intersection() — вжух — и всё готово!

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

UPD. После подсказки @rusorrow о том, что .keys() не нужен при создании множества, я подумал, что можно было бы сделать ещё короче: с помощью map — мне кажется, что так даже лучше — это, пожалуй, тот случай, когда map упрощает код по сравнению со списочными включениями.

In [7]:
set.intersection(*map(set, contracts))
Out[7]:
{'currency',
 'customer',
 'fileVersion',
 'fz',
 'id',
 'loadId',
 'mongo_id',
 'number',
 'price',
 'products',
 'publishDate',
 'regNum',
 'regionCode',
 'versionNumber'}

Комментарии