使用Scipy KDTree实现二维范围搜索时的匹配错误问题

背景信息 #

• 单位正方形上的二叉空间划分（类似k-d树单元格），以形状为(n,4)的numpy数组表示，n代表轴对齐、不重叠的边界框数量，每个边界框格式为[min_x, max_x, min_y, max_y]。
• 单位正方形内有p个采样点，以形状为(p,2)的数组表示。

问题描述 #

需要为每个采样点找到其所在边界框在数组中的索引，即执行正交范围搜索确定每个采样点所属的边界框。

尝试的解决方案 #

计算每个边界框的质心，基于质心构建k-d树，查询每个采样点的最近邻质心来匹配对应边界框。但使用scipy.spatial.KDTree得到错误结果，调整可选参数也无效。

问题详情 #

示例代码如下：

import numpy as np
from scipy.spatial import KDTree

# bbox coordinates as (min_x, max_x, min_y, max_y)
bbox = np.array([[0.75,       1.,         0.44,       1.        ],
        [0.,         0.1097561,  0.,         0.54666666],
        [0.1097561,  0.75,       0.,         0.54666666],
        [0.,         0.24264705, 0.54666666, 1.        ],
        [0.24264705, 0.75,       0.54666666, 1.        ],
        [0.75,       0.84090909, 0.,         0.44      ],
        [0.84090909, 1.,         0.,         0.44      ]])

# sample points
samples = np.array([[0.        , 0.        ],
                    [0.5       , 0.33333333],
                    [0.25      , 0.66666667],
                    [0.75      , 0.11111111],
                    [0.125     , 0.44444444]])


def find_bbox(bounding_boxes, samples):
    '''docstring'''
    
    # Calculate the centroid of each bounding box as (centroid_x, centroid_y).
    centroids = np.vstack((np.mean(bounding_boxes[:,0:2], axis=1),
                            np.mean(bounding_boxes[:,2:4], axis=1))).T
    # Build a scipy.spacial.KDTree from the bounding box centroids.
    tree = KDTree(centroids, leafsize=1)
    # For each XY-coordinate in (samples), find the nearest neighbour among
    # (centroids). 
    distances, indices = tree.query(samples, k=1)

    return indices


# Expected result would be: [1, 2, 4, 2, 2] or [1, 2, 4, 5, 2], instead I get
# [1, 2, 3, 5, 1]
print(find_bbox(bbox, samples))

运行后得到结果[1, 2, 3, 5, 1]，但预期结果应为[1, 2, 4, 2, 2]或[1, 2, 4, 5, 2]。

补充说明 #

意识到问题可能出在距离度量上：scipy.spatial.KDTree默认使用欧氏距离，当前结果符合该度量逻辑，但不符合需求。

解决方案 #

核心问题 #

用质心最近邻匹配的思路本身存在缺陷：一个点可能离某个边界框的质心更近，但实际落在另一个边界框内，尤其是当边界框大小差异较大时。必须直接判断点是否在边界框内，而非依赖质心距离。

基础实现方法 #

直接向量化判断每个点是否满足边界框的范围条件，提取对应索引：

import numpy as np

def find_bbox_correct(bounding_boxes, samples):
    # 扩展维度实现广播：(p,1,4) 与 (n,4) 匹配
    samples_expanded = samples[:, np.newaxis, :]
    # 判断每个点是否在每个bbox内：x在[min_x, max_x]且y在[min_y, max_y]
    in_bbox = (
        (samples_expanded[..., 0] >= bounding_boxes[:, 0]) &
        (samples_expanded[..., 0] <= bounding_boxes[:, 1]) &
        (samples_expanded[..., 1] >= bounding_boxes[:, 2]) &
        (samples_expanded[..., 1] <= bounding_boxes[:, 3])
    )
    # 边界框不重叠，每个点仅匹配一个bbox，提取对应索引
    indices = np.argmax(in_bbox, axis=1)
    return indices

# 测试
print(find_bbox_correct(bbox, samples))
# 输出：[1 2 4 5 2]，符合预期

优化方案（针对大量边界框/采样点） #

如果边界框和采样点数量都很大，上述O(p*n)复杂度的方法效率较低，可结合空间索引减少判断次数：

from scipy.spatial import cKDTree

def find_bbox_optimized(bounding_boxes, samples):
    centroids = np.column_stack((
        np.mean(bounding_boxes[:, :2], axis=1),
        np.mean(bounding_boxes[:, 2:], axis=1)
    ))
    tree = cKDTree(centroids)
    # 先找每个点周围的k个候选质心（k值可根据实际场景调整）
    _, candidate_indices = tree.query(samples, k=3)
    
    indices = np.zeros(len(samples), dtype=int)
    for i, (point, candidates) in enumerate(zip(samples, candidate_indices)):
        # 在候选中筛选实际包含点的bbox
        for idx in candidates:
            bbox = bounding_boxes[idx]
            if (bbox[0] <= point[0] <= bbox[1]) and (bbox[2] <= point[1] <= bbox[3]):
                indices[i] = idx
                break
    return indices

print(find_bbox_optimized(bbox, samples))
# 输出：[1 2 4 5 2]

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