python - 查找二维数组中接近元素的索引

Question

我有 2 个 nd 数组，其中每一行都是一个 3D 点，一个数组比另一个大得多。
即

array([[1., 2., 3.],
       [2.01, 5., 1.],
       [3., 3., 4.],
       [1., 4., 1.],
       [3., 6., 7.01]])

array([[3.02, 3.01, 4.0],
      [1.01, 1.99, 3.01],
      [2.98, 6.01, 7.01]])

而且我知道第二个数组中的每个点对应于第一个数组中的一个点。
我想获得通信索引列表，
即对于这个例子，它将是

 array([2,0,4])

由于第二个数组中的第一个点与第一个数组中的第三个点相似，第二个数组中的第二个点与第一个数组中的第一个点相似，依此类推。

Answer 1

您可以使用 KDTree 高效地做到这一点。

import numpy as np
from scipy.spatial import KDTree

x = np.array([[1., 2., 3.],
       [2.01, 5., 1.],
       [3., 3., 4.],
       [1., 4., 1.],
       [3., 6., 7.01]])

y = np.array([[1.01, 1.99, 3.01],
       [3.02, 3.01, 4.0],
       [2.98, 6.01, 7.01]])

result = KDTree(x).query(y)[1]

# In [16]: result                                                        
# Out[16]: array([0, 2, 4])

感谢 Divakar 指出 scipy 还提供了 KDTree 的 C 实现，称为 cKDTree。对于以下基准测试，它的速度提高了 10 倍:

x = np.random.rand(100_000, 3)
y = np.random.rand(100, 3)

def benchmark(TreeClass):
    return TreeClass(x).query(y)[1]

In [23]: %timeit q.benchmark(KDTree)                                   
322 ms ± 7.76 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [24]: %timeit q.benchmark(cKDTree)                                  
36.5 ms ± 763 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)