如何在KD树中实现范围搜索

Question

我已经构建了一个 d 维 KD-树。我想在这个树上执行区域搜索。维基百科 中提到了 KD-树中的范围搜索，但没有详细讲述其实现/算法。请问有人能帮我吗？即使不是对于任意的 d，也至少希望得到对于 d = 2 和 d = 3 的帮助。谢谢！

- Ankit Kumar

---

我找到了这些笔记http://www.cs.utah.edu/~lifeifei/cis5930/kdtree.pdf。它有关于d=1和d=2情况的伪代码。 - 9mat

@9mat 非常感谢。该算法可以推广到不同的 d 值，循环遍历节点/轴。 - Ankit Kumar

@9mat，你应该写一个回答，包括链接和内容的简要描述。我认为你指向了解决这个问题的最佳在线资源。 - m.raynal

Answer 1

kd-tree有多种变体，我使用的变体具有以下规格：

1. 每个内部节点最多只有两个节点。
2. 每个叶节点最多可以拥有maxCapacity个点。
3. 没有任何内部节点存储任何点。

注：还有版本是每个节点（无论是内部节点还是叶节点）都仅存储一个点。下面的算法也适用于这些版本。主要区别在于buildTree。

我大约两年前写了一个算法，感谢@9mat指向的资源。

假设任务是查找位于给定超矩形中（d维）的点数。这项任务也可以列出所有点或列出满足某些其他条件的范围内的所有点，但对我的代码进行直接更改即可。

定义基本节点类如下：

template <typename T> class kdNode{
    public: kdNode(){}
    virtual long rangeQuery(const T* q_min, const T* q_max) const{ return 0; }
};

那么，一个内部节点（非叶节点）可能长这样：

class internalNode:public kdNode<T>{
    const kdNode<T> *left = nullptr, *right = nullptr; // left and right sub trees
    int axis; // the axis on which split of points is being done
    T value; // the value based on which points are being split

    public: internalNode(){}

    void buildTree(...){
        // builds the tree recursively
    }

    // returns the number of points in this sub tree that lie inside the hyper rectangle formed by q_min and q_max
    int rangeQuery(const T* q_min, const T* q_max) const{
        // num of points that satisfy range query conditions
        int rangeCount = 0;

        // check for left node
        if(q_min[axis] <= value) {
            rangeCount += left->rangeQuery(q_min, q_max);
        }
        // check for right node
        if(q_max[axis] >= value) {
            rangeCount += right->rangeQuery(q_min, q_max);
        }

        return rangeCount;
    }
};

最终，叶节点应如下所示：

class leaf:public kdNode<T>{
    // maxCapacity is a hyper - param, the max num of points you allow a node to hold
    array<T, d> points[maxCapacity];
    int keyCount = 0; // this is the actual num of points in this leaf (keyCount <= maxCapacity)

    public: leaf(){}

    public: void addPoint(const T* p){
        // add a point p to the leaf node
    }

    // check if points[index] lies inside the hyper rectangle formed by q_min and q_max
    inline bool containsPoint(const int index, const T* q_min, const T* q_max) const{
        for (int i=0; i<d; i++) {
            if (points[index][i] > q_max[i] || points[index][i] < q_min[i]) {
                return false;
            }
        }
        return true;
    }

    // returns number of points in this leaf node that lie inside the hyper rectangle formed by q_min and q_max
    int rangeQuery(const T* q_min, const T* q_max) const{
        // num of points that satisfy range query conditions
        int rangeCount = 0;
        for(int i=0; i < this->keyCount; i++) {
            if(containsPoint(i, q_min, q_max)) {
                rangeCount++;
            }
        }
        return rangeCount;
    }
};

在叶子节点的范围查询代码中，可以在“线性搜索”内部进行“二分查找”。由于点将沿着轴axis排序，因此您可以使用q_min和q_max执行二分查找以查找l和r值，然后从l到r进行线性搜索，而不是从0到keyCount-1（当然，在最坏的情况下这没有帮助，但在实践中，特别是如果您具有相当高的容量值，这可能会有所帮助）。

- Ankit Kumar

---

你的代码有一个错误。假设根节点的值为50，因为这是中位数分割发生的地方。如果我们查询范围为15到30，它将返回0，因为根节点的值不在范围内，因此没有访问任何子节点。 - clankill3r

为了解决这个问题，你可以遍历到第一个分裂发生的节点。 另一件事是，有三种选项：完全在外面、相交和完全包含。当节点完全包含时，很好的一点是整个分支可以添加，而不需要检查值是否在范围内。 - clankill3r

希望我可以提供更详细的答案，问题是我仍在努力理解这一切。 - clankill3r

Answer 2

这是我的KD树解决方案，其中每个节点存储点（因此不仅仅是叶子节点）。 （请注意，适应仅在叶子节点中存储点的情况非常容易）。

我省略了一些优化，并将它们解释在最后，以减少解决方案的复杂性。 get_range函数在末尾具有可变参数，并且可以像以下方式调用， x1，y1，x2，y2或 x1，y1，z1，x2，y2，z2等。 首先给出范围的低值，然后是高值。 （您可以使用任意多个维度）。

static public <T> void get_range(K_D_Tree<T> tree, List<T> result, float... range) {
    if (tree.root == null) return;
   
    float[] node_region = new float[tree.DIMENSIONS * 2];
    for (int i = 0; i < tree.DIMENSIONS; i++) {
        node_region[i] = -Float.MAX_VALUE;
        node_region[i+tree.DIMENSIONS] = Float.MAX_VALUE;
    }
    
    _get_range(tree, result, tree.root, node_region, 0, range);
}

node_region 表示节点所在的区域，我们从尽可能大的区域开始。因为我们不知道这是否是我们正在处理的区域。

这里是递归 _get_range 的实现：

static public <T> void _get_range(K_D_Tree<T> tree, List<T> result, K_D_Tree_Node<T> node, float[] node_region, int dimension, float[] target_region) {
    if (dimension == tree.DIMENSIONS) dimension = 0;

    if (_contains_region(tree, node_region, target_region)) {
        _add_whole_branch(node, result);
    }
    else {
        float value = _value(tree, dimension, node);

        if (node.left != null) {
            float[] node_region_left  = new float[tree.DIMENSIONS*2];
            System.arraycopy(node_region, 0, node_region_left, 0, node_region.length);
            node_region_left[dimension + tree.DIMENSIONS] = value;

            if (_intersects_region(tree, node_region_left, target_region)){
                _get_range(tree, result, node.left, node_region_left, dimension+1, target_region);
            }
        }

        if (node.right != null) {
            float[] node_region_right = new float[tree.DIMENSIONS*2];
            System.arraycopy(node_region, 0, node_region_right, 0, node_region.length);
            node_region_right[dimension] = value;

            if (_intersects_region(tree, node_region_right, target_region)){
                _get_range(tree, result, node.right, node_region_right, dimension+1, target_region);
            }
        }        

        if (_region_contains_node(tree, target_region, node)) {
            result.add(node.point);
        }
    }
}

另一个答案没有提供的重要一点是这部分内容：

if (_contains_region(tree, node_region, target_region)) {
    _add_whole_branch(node, result);
}

使用KD-Tree进行范围搜索时，节点的区域有三种选项：

1. 完全在外部
2. 相交
3. 完全包含

一旦确定一个区域完全包含，则可以添加整个分支而无需进行任何维度检查。 为了更清楚地说明，这里是_add_whole_branch：

static public <T> void _add_whole_branch(K_D_Tree_Node<T> node, List<T> result) {
    result.add(node.point);
    if (node.left != null)  _add_whole_branch(node.left, result);
    if (node.right != null) _add_whole_branch(node.right, result);
}

在这张图片中，所有大的白点都是使用_add_whole_branch添加的，而仅对于红点需要检查所有维度。

优化

1) 与其从根节点开始_get_range函数，不如找到分裂节点。这是第一个其点在查询范围内的节点。要找到分裂节点，您仍需要从根节点开始，但计算会更便宜（因为您只需向左或向右移动）。

2) 现在我创建了float[] node_region_left和float[] node_region_right，由于这发生在递归函数中，可能会导致相当多的数组。但是，您可以重复使用左侧的数组来处理右侧的数组。出于清晰起见，我在此示例中没有这样做。 我还可以想象将区域大小存储在节点中，但这需要更多的内存，并可能导致许多缓存未命中。

- clankill3r