如何使用numpy配对（x，y）对

Question

我有2个2D NumPy数组，每个数组由〜300,000（x，y）对组成。给定数组A中的^第 i （x，y）对，我需要找到数组B中的^第 j （x，y）对，使得（[x _i -x _j ] ² + [y _i -y _j ] ²）^½（两个（x，y）对之间的距离）最小。

我目前正在做的是argmin这样的：

thickness = []
for i in range(len(A)):
    xi = A[i][0]
    yi = A[i][1]
    idx = (np.sqrt(np.power(B[:, 0] - xi, 2) + np.power(B[:, 1] - yi, 2))).argmin()
    thickness.append([(xi + B[idx][0]) / 2, (yi + B[idx][1]) / 2,
                      A[i][2] + B[idx][2]])

有没有更快的方法可以做到这一点？

Answer 1

我们可以使用Cython-powered kd-tree快速最近邻查找来获取最近邻，从而实现我们想要的输出，如下所示 -

from scipy.spatial import cKDTree

idx = cKDTree(B[:,:2]).query(A[:,:2], k=1)[1]
thickness = [(A[:,0] + B[idx,0]) / 2, (A[:,1] + B[idx,1]) / 2, A[:,2] + B[idx,2]]