python - 查找数组中的最近点 - KDTree 的倒数

Question

我有一个非常大的 ndarray A 和一个排序的点列表 k(一个小列表，大约 30 个点)。

对于 A 的每个元素，我想确定点列表 k 中最近的元素以及索引。所以像这样:

>>> A = np.asarray([3, 4, 5, 6])
>>> k = np.asarray([4.1, 3])
>>> values, indices
[3, 4.1, 4.1, 4.1], [1, 0, 0, 0]

现在，问题是 A 非常非常大。所以我不能做一些低效的事情，比如给 A 添加一个维度，将 abs 差值取为 k，然后取每一列的最小值。

现在我一直在使用 np.searchsorted，如这里的第二个答案所示:Find nearest value in numpy array但即使这样也太慢了。这是我使用的代码(修改为使用多个值):

def find_nearest(A,k):

    indicesClosest = np.searchsorted(k, A)
    flagToReduce = indicesClosest==k.shape[0]
    modifiedIndicesToAvoidOutOfBoundsException = indicesClosest.copy()
    modifiedIndicesToAvoidOutOfBoundsException[flagToReduce] -= 1
    flagToReduce = np.logical_or(flagToReduce,
                     np.abs(A-k[indicesClosest-1]) <
                     np.abs(A - k[modifiedIndicesToAvoidOutOfBoundsException]))
    flagToReduce = np.logical_and(indicesClosest > 0, flagToReduce)
    indicesClosest[flagToReduce] -= 1
    valuesClosest = k[indicesClosest]
    return valuesClosest, indicesClosest

然后我想到了使用 scipy.spatial.KDTree:

>>> d = scipy.spatial.KDTree(k)
>>> d.query(A)

事实证明，这比搜索排序的解决方案慢得多。

另一方面，数组A总是一样的，只有k变了。因此，最好在 A 上使用一些辅助结构(如“逆向 KDTree”)，然后在小数组 k 上查询结果。

有这样的东西吗？

编辑

目前我正在使用 np.searchsorted 的变体，它需要对数组 A 进行排序。我们可以提前做这个作为预处理步骤，但是我们仍然需要在计算索引后恢复原来的顺序。此变体的速度大约是上述变体的两倍。

A = np.random.random(3000000)
k = np.random.random(30)

indices_sort = np.argsort(A)
sortedA = A[indices_sort]

inv_indices_sort = np.argsort(indices_sort)
k.sort()


def find_nearest(sortedA, k):
    midpoints = k[:-1] + np.diff(k)/2
    idx_aux = np.searchsorted(sortedA, midpoints)
    idx = []
    count = 0
    final_indices = np.zeros(sortedA.shape, dtype=int)
    old_obj = None
    for obj in idx_aux:
        if obj != old_obj:
            idx.append((obj, count))
            old_obj = obj
        count += 1
    old_idx = 0
    for idx_A, idx_k in idx:
        final_indices[old_idx:idx_A] = idx_k
        old_idx = idx_A
    final_indices[old_idx:] = len(k)-1

    indicesClosest = final_indices[inv_indices_sort] #<- this takes 90% of the time
    return k[indicesClosest], indicesClosest

花费这么多时间的那条线是将索引恢复到原来顺序的那条线。

Answer 1

更新:

内置函数 numpy.digitize实际上可以完全满足您的需求。只需要一个小技巧:digitize 将值分配给 bins。我们可以将 k 转换为 bin，方法是对数组进行排序并将 bin 边界正好设置在相邻元素之间的中间。

import numpy as np

A = np.asarray([3, 4, 5, 6])
k = np.asarray([4.1, 3, 1])  # added another value to show that sorting/binning works

ki = np.argsort(k)
ks = k[ki]

i = np.digitize(A, (ks[:-1] + ks[1:]) / 2)

indices = ki[i]
values = ks[i]

print(values, indices)
# [ 3.   4.1  4.1  4.1] [1 0 0 0]

---

旧答案:

我会采用蛮力方法对 k 中的每个元素执行一次矢量化遍历 A 并更新当前元素改进近似值的那些位置。

import numpy as np

A = np.asarray([3, 4, 5, 6])
k = np.asarray([4.1, 3])

err = np.zeros_like(A) + np.inf  # keep track of error over passes

values = np.empty_like(A, dtype=k.dtype)
indices = np.empty_like(A, dtype=int)

for i, v in enumerate(k):
    d = np.abs(A - v)
    mask = d < err  # only update where v is closer to A
    values[mask] = v
    indices[mask] = i
    err[mask] = d[mask]

print(values, indices)
# [ 3.   4.1  4.1  4.1] [1 0 0 0]

此方法需要三个与 A 大小相同的临时变量，因此如果没有足够的可用内存，它将失败。