Словарь словарей множеств из pandas DataFrame для NetworkX

Покажем, как из pandas DataFrame быстро получить словарь словарей множеств с ненулевыми весами для построения графа в NetworkX: понятный код и примеры.

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

Постановка задачи

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

payload = {'I': ['A', 'B', 'C', 'D'], 'X': [1, 0, 3, 1], 'Y': [0, 1, 2, 1], 'Z': [1, 0, 0, 0], 'W': [3, 2, 0, 0]}
frame = pd.DataFrame(data=payload, columns=['I','X','Y','Z','W'])
frame.set_index('I', inplace=True, drop=True)

Желаемая структура — это словарь, который отображает узел в его ненулевые атрибуты и их веса, где каждый вес представлен множеством:

{'A': {'X': {1}, 'Z': {1}, 'W': {3}}, 'B': {'Y': {1}, 'W': {2}}, 'C': {'X': {3}, 'Y': {2}}, 'D': {'Y': {1}, 'X': {1}}}

Что здесь происходит

Каждая строка кодирует узел, каждый столбец — потенциальный атрибут ребра, а значение в ячейке — его вес. Нужно пропускать нули и переносить только существующие рёбра, оборачивая каждый скалярный вес в множество из одного элемента. Такой формат удобен для построения и анализа графа в Networkx, но собирать его по ячейкам плохо масштабируется.

Подходы к решению

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

edge_map = {
    src: {attr: {val} for attr, val in attrs.items() if val != 0}
    for src, attrs in frame.iterrows()
}

Если хотите ещё сильнее снизить накладные расходы на итерацию в pandas, сначала преобразуйте DataFrame в обычный словарь, а затем соберите вложенную структуру. Так вся трансформация выполняется в нативных структурах данных Python:

row_dict = frame.to_dict(orient='index')
edge_map = {
    src: {k: {v} for k, v in inner.items() if v != 0}
    for src, inner in row_dict.items()
}

Это можно сжать в одно включение без промежуточной переменной:

edge_map = {
    node: {col: {val} for col, val in cols.items() if val}
    for node, cols in frame.to_dict(orient='index').items()
}

Почему это важно

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

Итоги

Когда из DataFrame с весами рёбер нужен словарь словарей множеств, опирайтесь на построчное преобразование в словарь и одно фильтрующее включение. Если хотите оставаться в экосистеме pandas, итерация по строкам для этой конкретной задачи всё ещё достаточно эффективна и хорошо читается. А если возможности DataFrame за пределами этого шага не нужны, работайте сразу с обычным словарём: так проще и, как правило, быстрее обрабатывать.