How to pair (x,y) pairs using numpy

Question

I have 2, 2D NumPy arrays consisting of ~300,000 (x,y) pairs each. Given the i^th (x, y) pair in array A, I need to find the corresponding j^th (x,y) pair in array B such that ([x_i - x_j]² + [y_i - y_j]²)^½, the distance between the two (x,y) pairs, is minimized.

What I'm currently doing is argmin like so:

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]])

Is there a faster way to accomplish this?

Answer 1

We can use Cython-powered kd-tree for quick nearest-neighbor lookup to get nearest neighbour and hence achieve our desired output, like so -

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]]