基于Numpy距离矩阵高效生成JSON格式的用户最近匹配摘要

我来帮你解决这个问题——处理Numpy距离矩阵并高效提取每个用户的最近n个匹配，最终输出JSON格式的结果，这是我常用的一套高效实现方案：

核心思路 #

1. 遍历距离矩阵中的每个用户，排除自身（对角线的inf值）
2. 用高效的分区算法找到每个用户对应的最小n个距离（比全排序更快）
3. 将结果整理为嵌套字典结构，最后序列化为JSON

完整代码实现 #

首先导入依赖库：

import numpy as np
import json

定义示例距离矩阵（补全你给出的5×5矩阵）：

# 示例5×5距离矩阵
dist_matrix = np.array([
    [np.inf, 0.30330249, 0.41690763, 0.11468943, 0.27026611],
    [0.30330249, np.inf, 0.52100912, 0.22100123, 0.18956789],
    [0.41690763, 0.52100912, np.inf, 0.35678901, 0.29876543],
    [0.11468943, 0.22100123, 0.35678901, np.inf, 0.15678902],
    [0.27026611, 0.18956789, 0.29876543, 0.15678902, np.inf]
])

实现核心处理函数：

def get_top_n_matches(dist_matrix, n):
    num_users = dist_matrix.shape[0]
    result = []
    
    for user_idx in range(num_users):
        # 获取当前用户的所有距离，过滤掉自身的inf值
        distances = dist_matrix[user_idx]
        other_user_indices = np.where(distances != np.inf)[0]
        target_distances = distances[other_user_indices]
        
        # 用argpartition高效定位最小的n个元素（时间复杂度O(N)，远快于全排序的O(N log N)）
        partition_pos = np.argpartition(target_distances, n)[:n]
        candidate_indices = other_user_indices[partition_pos]
        candidate_distances = target_distances[partition_pos]
        
        # 可选：对候选结果按距离从小到大排序（如果需要严格有序的匹配列表）
        sorted_order = np.argsort(candidate_distances)
        sorted_indices = candidate_indices[sorted_order]
        sorted_distances = candidate_distances[sorted_order]
        
        # 构建当前用户的匹配详情字典
        user_match_info = {
            "user_id": int(user_idx),
            "top_n_matches": [
                {"matched_user_id": int(match_id), "distance": round(float(dist), 8)}
                for match_id, dist in zip(sorted_indices, sorted_distances)
            ]
        }
        result.append(user_match_info)
    
    return result

调用函数并输出JSON：

# 设置要提取的最近匹配数量n
n = 2
# 获取所有用户的最近n个匹配
top_matches = get_top_n_matches(dist_matrix, n)
# 序列化为格式化的JSON字符串
json_result = json.dumps(top_matches, indent=4)
print(json_result)

关键细节说明 #

• 效率优化：使用np.argpartition而非np.argsort，当用户数量很大（比如上万级）时，这种方法能大幅减少计算时间。如果不需要严格排序的匹配列表，可以直接去掉最后的argsort步骤，进一步提升效率。
• 排除自身：通过np.where(distances != np.inf)过滤掉对角线的自身距离，避免出现用户匹配自己的情况。
• 数据兼容性：将Numpy的数值类型转换为Python原生的int/float，确保JSON序列化时不会出错；同时对距离值做了保留8位小数的处理，让输出更整洁。

示例输出 #

运行上述代码后，会得到如下格式化的JSON结果：

[
    {
        "user_id": 0,
        "top_n_matches": [
            {
                "matched_user_id": 3,
                "distance": 0.11468943
            },
            {
                "matched_user_id": 4,
                "distance": 0.27026611
            }
        ]
    },
    {
        "user_id": 1,
        "top_n_matches": [
            {
                "matched_user_id": 4,
                "distance": 0.18956789
            },
            {
                "matched_user_id": 3,
                "distance": 0.22100123
            }
        ]
    },
    {
        "user_id": 2,
        "top_n_matches": [
            {
                "matched_user_id": 4,
                "distance": 0.29876543
            },
            {
                "matched_user_id": 3,
                "distance": 0.35678901
            }
        ]
    },
    {
        "user_id": 3,
        "top_n_matches": [
            {
                "matched_user_id": 4,
                "distance": 0.15678902
            },
            {
                "matched_user_id": 0,
                "distance": 0.11468943
            }
        ]
    },
    {
        "user_id": 4,
        "top_n_matches": [
            {
                "matched_user_id": 3,
                "distance": 0.15678902
            },
            {
                "matched_user_id": 1,
                "distance": 0.18956789
            }
        ]
    }
]

内容的提问来源于stack exchange，提问作者sjw