python - 从 numpy 距离矩阵高效生成 JSON 摘要

Question

我有一个存储为二维 numpy 数组的距离矩阵。我正在寻找一种有效的方法来为人口中的每个用户提取包含最接近的 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，我希望每个用户前 25 个而不是前 2 个(适当大小的示例可以通过在上面的代码中将 n 设置为 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})