Эффективно сгенерировать JSON-сводку из матрицы пустых расстояний

Question

У меня есть матрица расстояний, хранящаяся в виде двумерного массива. Я ищу эффективный способ извлечь сводку, содержащую детали ближайших n совпадений, для каждого пользователя в популяции. Эта сводка в конечном итоге будет использоваться как JSON, поэтому я хочу, чтобы она была в виде вложенного списка / словаря (пример вывода чуть ниже).

Следующий минимальный пример (матрица расстояний 5 x 5) демонстрирует, с чем я работаю:

[[       inf 0.30330249 0.41690763 0.11468943 0.27026611]
 [0.30330249        inf 0.72021012 0.41799192 0.5735686 ]
 [0.41690763 0.72021012        inf 0.3022182  0.14664152]
 [0.11468943 0.41799192 0.3022182         inf 0.15557668]
 [0.27026611 0.5735686  0.14664152 0.15557668        inf]]

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

import numpy as np
from scipy.spatial.distance import squareform, pdist

n = 5  # Population size
np.random.seed(1)
users = ['User {}'.format(i) for i in range(1, n+1)]
dm = squareform(pdist(np.random.random((n, 1))))
np.fill_diagonal(dm, np.inf)

Допустим, мы хотим найти 2 ближайших совпадения для каждого пользователя. Посмотрев на матрицу расстояний, мы видим, что для «Пользователя 1» их самые близкие совпадения - «Пользователь 4» (0.11468943), а затем «Пользователь 5» (0.27026611). Мой желаемый вывод выглядит следующим образом:

{
    "User 1": [
        {
            "Main": "User 1",
            "Other": "User 4",
            "Distance": 0.11468943207073423
        },
        {
            "Main": "User 1",
            "Other": "User 5",
            "Distance": 0.27026611388546096
        }
    ],
    "User 2": [
        # redacted
    ],
    "User 3": [
        # redacted
    ],
    "User 4": [
        # redacted
    ],
    "User 5": [
        {
            "Main": "User 5",
            "Other": "User 3",
            "Distance": 0.14664151599976816
        },
        {
            "Main": "User 5",
            "Other": "User 4",
            "Distance": 0.15557668181472672
        }
    ]
}

(я понимаю, что приведенные выше клавиши "Main" немного избыточны, я включил их, чтобы облегчить работу с данными на внешнем интерфейсе)

Мне удалось достичь желаемых результатов, используя следующий код:

import pandas as pd

n_per_user = 2  # Number of closest users to find per user

# Get row-wise indices of n smallest distances
indices = np.argpartition(dm, range(n_per_user), axis=1)[:, :n_per_user]

# Each of these comprehensions is for one column of the DataFrame which will be built shortly
users_main = (user for user in users for i in range(n_per_user))
users_other = (users[i] for i in indices.flatten())
distances = (dm[i, j] for i, row in enumerate(indices) for j in row)

# Construct the DataFrame
df = pd.DataFrame(list(zip(users_main, users_other, distances)), columns=['Main', 'Other', 'Distance'])

#      Main   Other  Distance
# 0  User 1  User 4  0.114689
# 1  User 1  User 5  0.270266
# 2  User 2  User 1  0.303302
# 3  User 2  User 4  0.417992
# 4  User 3  User 5  0.146642
# 5  User 3  User 4  0.302218
# 6  User 4  User 1  0.114689
# 7  User 4  User 5  0.155577
# 8  User 5  User 3  0.146642
# 9  User 5  User 4  0.155577

results = {x: y.to_dict('records') for x, y in df.groupby('Main', sort=False)}

Это хорошо работает для крошечных наборов данных, подобных этому, но мой реальный dm - это 10k x 10k, а не 5 x 5, и я хочу, чтобы top 25 приходилось на пользователя, а не top 2 (пример соответствующего размера можно сгенерировать, установив * От 1023 * до 10000 и n_per_user до 25 в приведенном выше коде).

Вся программа в ее текущем состоянии выполняется на моем компьютере примерно за 10 секунд, причем самый последний шаг (преобразование DataFrame во вложенный словарь) занимает более половины этого времени. Учитывая, что я хотел бы, чтобы эти шаги выполнялись очень часто в конечном приложении, я ищу более эффективное решение. Я понимаю, что мог бы просто попросить помощи на этом последнем шаге, поскольку именно он является причиной узкого места, но я подозреваю, что могут быть более эффективные решения, которые вообще обходят необходимость создания DataFrame, поэтому я включил так много контекста.

Answer 1

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

from collections import defaultdict

results = defaultdict(list)
for main, other, distance in zip(users_main, users_other, distances):
    results[main].append({"Main": main, "Other": other, "Distance": distance})