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

【问题标题】：Find closest k points for every point in row of numpy array为 numpy 数组行中的每个点查找最近的 k 点
【发布时间】：2017-03-17 11:58:30
【问题描述】：

我有一个 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 函数可以实现这一点？

【问题讨论】：

“最接近”是什么意思？按价值最接近？什么是“点”？
所以说行 r = [1,10,11,15,18,16,19,18]。对于 r 中的每个元素，我想找到 r 中具有最接近我们正在查看的元素的值的其他 k 个元素。所以我们看 1 并找到最接近它的 k 个数字。然后我们查看 10 并找到最接近它的 k 个数字，然后......然后 18 并找到最接近它的 k 个数字。
是的，它是一个预测矩阵，所以行是人，列是项目
那么，对于每一行，您希望得到一个 1000x5 的数组作为结果吗？
既然您的指标如此简单，您就不能只使用sort（或argsort）您的行吗？这会立即将候选最近邻居的数量减少到每个点 2k，其中 k 是您示例中的 5。

标签： python sorting numpy knn

【解决方案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)

【讨论】：

干得好，也打败我吧。我要补充一点，与在循环中根据元素和行之间的差异对向量执行 np.argpartition 的方法相比，这要快 6 到 7 倍。（~ 3 秒诗句~ 18 秒）。我认为更高版本的 scipy 有一个 n_jobs 参数用于 tree.query 函数，用于跨 CPU 内核进行并行处理。我的版本没有该参数，但这也可能会提高性能。

【解决方案2】：

我不太确定您希望最终结果如何。但这绝对可以满足您的需求。

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

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

x1, x2 = np.triu_indices(X.shape[1], 1)

生成所有距离的数组

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

为每一行找到最接近的 5 个

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

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

【讨论】：

【解决方案3】：

这是一个argsorting 解决方案，它努力利用简单的指标：

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)

输出：

# 0.00150537172454
# 0.4567880852999224

【讨论】：