从numpy距离矩阵中有效地生成JSON摘要

时间:2018-05-02 16:05:41

标签: python python-3.x pandas numpy

我将距离矩阵存储为二维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,这就是为什么我包含了这么多的上下文。 / p>

1 个答案:

答案 0 :(得分: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})