为数字数组行中的每个点找到最近的 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 树可以为此工作吗？我尝试了一种类似于

在 x 和 y 坐标的 numpy 数组中查找最近点的索引

但是，我无法让它工作。是否有一些我不知道的 numpy 功能可以完成此操作？

Answer 1

构造一个 kdtree，每行数据都有scipy.spatial.cKDTree。

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

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

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]给出最接近对的第二个位置。

Answer 3

这是一个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)

输出：

# 0.00150537172454
# 0.4567880852999224