使用KDTree匹配NumPy数组最近邻时结果错误的问题排查

我有两个NumPy数组：较小的整数数组a，较大的浮点数数组b，b中包含与a中部分整数接近的浮点数。目标是为a中的每个元素找到b里的最近值。

朴素for循环方法已得到正确结果：

a = np.array([35, 11, 48, 20, 13, 31, 49])
b = np.array([34.78, 34.8, 35.1, 34.99, 11.3, 10.7, 11.289, 18.78, 19.1, 20.05, 12.32, 12.87, 13.5, 31.03, 31.15, 29.87, 48.1, 48.5, 49.2])

for e in a:
    idx = np.abs(e - b).argsort()
    print(f"{e} has nearest match = {b[idx[0]]:.4f}")
# 输出：
# 35 has nearest match = 34.9900
# 11 has nearest match = 11.2890
# 48 has nearest match = 48.1000
# 20 has nearest match = 20.0500
# 13 has nearest match = 12.8700
# 31 has nearest match = 31.0300
# 49 has nearest match = 49.2000

实际场景中a.size=2040，b.size=1041901，尝试用KDTree实现时出现问题：

# 错误的KDTree实现
from scipy.spatial import KDTree

kd_tree = KDTree(data = np.expand_dims(a, 1))
dist_nn, idx_nn = kd_tree.query(x = np.expand_dims(b, 1), k = [1])

print(dist_nn.shape, idx_nn.shape)
# ((19, 1), (19, 1))
print(b[idx_nn])
# 输出：
# array([[10.7  ],
#        [10.7  ],
#        [10.7  ],
#        [11.289],
#        [11.289],
#        [11.289],
#        [11.3  ],
#        [11.3  ],
#        [11.3  ],
#        [12.32 ],
#        [12.32 ],
#        [12.32 ],
#        [12.87 ],
#        [12.87 ],
#        [12.87 ],
#        [12.87 ],
#        [13.5  ],
#        [13.5  ],
#        [18.78 ]])

遇到的问题：

• KDTree似乎未处理a中大于20的元素，如31、48、49等完全被遗漏；
• 多数最近邻结果与for循环的正确结果不符。

问题原因与解决方法 #

核心错误：KDTree的构建与查询对象搞反了 #

你的需求是为a中的每个元素找b里的最近值，但当前代码用a构建KDTree，再用b去查询——这实际是在找每个b元素在a中的最近整数，和需求完全相反，这就是结果错误的根本原因。

正确的KDTree实现 #

应该用较大的数组b构建KDTree，然后用较小的数组a去查询，这样就能高效得到每个a元素在b中的最近邻：

from scipy.spatial import KDTree

# 用b构建KDTree，需将一维数组扩展为二维（KDTree要求输入为2D数组）
kd_tree = KDTree(np.expand_dims(b, 1))
# 用a查询每个元素的最近1个邻点
dist_nn, idx_nn = kd_tree.query(np.expand_dims(a, 1), k=1)

# 输出匹配结果
for e, nearest_val in zip(a, b[idx_nn]):
    print(f"{e} has nearest match = {nearest_val:.4f}")

运行结果与朴素循环完全一致：

35 has nearest match = 34.9900
11 has nearest match = 11.2890
48 has nearest match = 48.1000
20 has nearest match = 20.0500
13 has nearest match = 12.8700
31 has nearest match = 31.0300
49 has nearest match = 49.2000

效率说明 #

当b规模远大于a时（如你的场景：b是百万级，a是两千级），用b构建一次KDTree，再对a做O(logN)的查询，总时间复杂度为O(M logN)（M为a的大小，N为b的大小），相比朴素循环的O(M*N)，效率提升极其显著。

内容的提问来源于stack exchange，提问作者Arun