为numpy数组行中的每个点查找最近的k点

Question

我有一个np数组，它的大小是1000x1000，其中每个元素都是实数。我想找到这个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

为数据的每一行构造一个带有^{}的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)

Answer 2

下面是一个argsort的解决方案，它努力利用简单的度量：

def nn(A, k):
    out = np.zeros((A.shape[0], A.shape[1] + 2*k), dtype=int)
    out[:, k:-k] = np.argsort(A, axis=-1)
    out[:, :k] = out[:, -k-1, None]
    out[:, -k:] = out[:, k, None]
    strd = stride_tricks.as_strided(
        out, strides=out.strides + (out.strides[-1],), shape=A.shape + (2*k+1,))
    delta = A[np.arange(A.shape[0])[:, None, None], strd]
    delta -= delta[..., k, None]
    delta = np.abs(delta)
    s = np.argpartition(delta,(0, k), axis = -1)[..., 1:k+1]
    inds = tuple(np.ogrid[:strd.shape[0], :strd.shape[1], :0][:2])
    res = np.empty(A.shape + (k,), dtype=int)
    res[np.arange(strd.shape[0])[:, None, None], out[:, k:-k, None],
        np.arange(k)[None, None, :]] = strd[inds + (s,)]
    return res

N = 1000
k = 5
r = 10

A = np.random.random((N, N))
# crude test
print(np.abs(A[np.arange(N)[:, None, None], res]-A[..., None]).mean())
# timings
print(timeit(lambda: nn(A, k), number=r) / r)

输出：

^{pr2}$

Answer 3

我不太确定你想要怎样的最终结果。但这绝对能满足你的需要。在

np.random.seed([3,1415])
X = np.random.rand(1000, 1000)

抓取上面的三角形索引来跟踪每行的每个点的组合

^{pr2}$

生成所有距离的数组

d = np.abs(X[:, x1] - X[:, x2])

每行找出最接近的5

tpos = np.argpartition(d, 5)[:, :5]

然后x1[tpos]给出最接近对中第一个点的行方向位置，而x2[tpos]给出最近对中的第二个位置。在