python - 为 numpy 数组的行中的每个点找到最近的 k 个点

Question

我有一个 np 数组 X，大小为 1000 x 1000，其中每个元素都是实数。我想为这个 np 数组的每一行中的每个点找到 5 个最近的点。这里的距离度量可以只是 abs(x-y)。我试过了

for i in range(X.shape[0]):
    knn = NearestNeighbors(n_neighbors=5)
    knn.fit(X[i])
    for j in range(X.shape[1])
        d = knn.kneighbors(X[i,j], return_distance=False)

但是，这对我不起作用，我不确定它的效率如何。有没有解决的办法？我见过很多用于比较向量的方法，但没有看到任何用于比较单个元素的方法。我知道我可以使用 for 循环并循环并找到最小的 k，但这在计算上会很昂贵。 KD 树可以解决这个问题吗？我试过类似

的方法

Finding index of nearest point in numpy arrays of x and y coordinates

但是，我无法让它工作。是否有一些我不知道的 numpy 函数可以实现此目的？

Answer 1

用 scipy.spatial.cKDTree 构造一个 kdtree对于每一行数据。

import numpy as np
import scipy.spatial


def nearest_neighbors(arr, k):
    k_lst = list(range(k + 2))[2:]  # [2,3]
    neighbors = []

    for row in arr:
        # stack the data so each element is in its own row
        data = np.vstack(row)
        # construct a kd-tree
        tree = scipy.spatial.cKDTree(data)
        # find k nearest neighbors for each element of data, squeezing out the zero result (the first nearest neighbor is always itself)
        dd, ii = tree.query(data, k=k_lst)
        # apply an index filter on data to get the nearest neighbor elements
        closest = data[ii].reshape(-1, k)
        neighbors.append(closest)
    return np.stack(neighbors)


N = 1000
k = 5
A = np.random.random((N, N))
nearest_neighbors(A, k)