如何在 numpy 中找到最近的邻居？

Question

There are two array u and v.

u.shape = (N,d) v.shape = (q,d)

I need to find, for every q, the nearest value's index for each d in u.

For example:

u = [[5,3],
     [3,4],
     [3,2],
     [8,7]] , shape (4,2)
v = [[1,3],
     [2,4]] , shape (2,2)

and I found many people said we can do that:

v = v.expand_dims(v,axis=1) # reshape to (2,1,2) for broadcast

result = np.argmin(abs(v-u),axis=1) # (u-v).shape = (2,4,2)

Of course it found the nearest value's index. But, when there are two nearest value. I need to take the "second" one's index.

In that case:

v-u = [[[-4,  0],
        [-2, -1],
        [-2,  1],
        [-7, -4]],

       [[-3,  1],
        [-1,  0],
        [-1,  2],
        [-6, -3]]])

along axis=1, there are two -2 in (uv)[0,:,0] and two -1 in (uv)[1,:,0] If we directly use:

result = np.argmin(abs(v-u),axis=1)

result will be:

array([[1, 0],
       [1, 1]], dtype=int64)

It returns the indices corresponding to the first occurrence but I need the second one, i,e

array([[2, 0],
       [2, 1]], dtype=int64)

Can anyone help? Thanks!

Answer 1

If there can be at most 2 minimal values, you can retrieve indices of the last minimum.

To do it:

• reverse abs(vu) along axis 1,
• compute argmin , getting a "reversed_index" (actually the index in the reversed array),
• map back to "original" indices using u.shape[0] - 1 - <reversed_index> formula (in your case of 4 rows, reversed index == 3 corresponds to original index == 0 )

The whole code is:

u.shape[0] - 1 - np.argmin(abs(v-u)[:,::-1,:],axis=1)

Other choice , when there can be more than 2 min values, is to write a specialized version of argmin , for an 1-D input array, returning the index of the second minimal value if there are more of them:

def argmin2(arr):
    ind = arr.argpartition(1)[:2]
    return ind[0] if arr[ind[0]] < arr[ind[1]] else ind[1]

and then apply it to abs(vu) along axis 1 :

np.apply_along_axis(argmin2, 1, abs(v-u))