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

Question

有两个数组u和v。

u.shape = (N，d( V.形状 = (q，d(

我需要为每个 q 找到 u 中每个 d 的最接近值的索引。

例如：

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

我发现很多人说我们可以做到这一点：

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)

当然，它找到了最接近的值的索引。但！当有两个最接近的值时，我需要取"第二个"的索引。

在这种情况下：

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

沿轴=1，有两个 -2 in (u-v([0，：，0] 和两个 -1 in (u-v([1，：，0] 如果我们直接使用：

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

结果将是：

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

它返回对应于第一次出现的索引，但我需要第二个，即

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

谁能帮忙？谢谢！

Answer 1

如果最多可以有 2 个最小值，则可以检索最后一个最小值。

要做到这一点：

• 沿轴 1 反向abs(v-u(，
• 计算参数，得到一个"reversed_index"(实际上是 反转阵列(，
• 使用U.shape[0] - 1 -公式映射回"原始"索引(在 4 行的情况下，反转索引 == 3对应于原始索引 == 0(

整个代码是：

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

当可以有超过 2 分钟的值时，其他选择是写入Argmin的专用版本，用于一维输入数组，返回第二个最小值的索引(如果有更多(：

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

然后沿轴1将其应用于ABS(V-U(：

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