NumPy二维数组非零元素到最近非零元素欧氏距离的高效Python实现

优化解法 #

我们可以先提取所有非零元素的坐标，再通过KDTree快速查询最近邻，时间复杂度为O(K log K)（K为非零元素个数），远优于暴力解法，代码也更简洁：

import numpy as np
from scipy.spatial import cKDTree

def get_distance_list(A, decimals=2):
    # 提取非零元素坐标，形状为(K, 2)，K为非零元素总数
    coords = np.column_stack(np.nonzero(A))
    k = coords.shape[0]
    if k < 2:
        return []
    # 构建KD树，查询每个点最近的2个邻点（第一个是自身，第二个为目标最近点）
    tree = cKDTree(coords)
    distances, _ = tree.query(coords, k=2)
    # 取第二个邻点的距离，保留指定小数位
    return distances[:, 1].round(decimals).tolist()

验证测试 #

用你给出的示例矩阵测试：

A = np.array([[0, 0, 0, 1, 0, 2, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 1, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0, 3, 0], [1, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 6, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0]])
print(get_distance_list(A))

输出结果和示例完全一致：[2.0, 2.0, 2.24, 2.24, 2.24, 2.83, 2.24, 2.24, 2.24]

优势说明 #

• 性能提升显著：当矩阵稀疏、非零元素较少时，速度比暴力解法高几个数量级，即使非零元素较多，KDTree的查询效率也远高于两两计算距离。
• 代码简洁易读：核心逻辑仅几行，完全符合Pythonic风格。
• 完全满足不变性规则：非零元素小于2时返回空列表，否则返回列表长度等于非零元素总数。

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