Kd-Trees

凝神长老

发布于 2020-05-07 21:10:23

7760

发布于 2020-05-07 21:10:23

文章被收录于专栏：凝神长老和他的朋友们凝神长老和他的朋友们

KD 树有许多应用，从对天文物体进行分类到计算机动画，再到加速神经网络，再到挖掘数据再到图像检索等。

下面以普林斯顿大学算法课第 5 次作业为例，长老向大家分享这种高效、神奇的数据结构。

题目链接：https://coursera.cs.princeton.edu/algs4/assignments/kdtree/specification.php

本次课程作业是编写一个数据结构，以表示单位正方形中的一组点，并支持高效的范围搜索（查找查询矩形中包含的所有点），以及高效的最近邻居搜索（找到最接近查询点的点）。

首先要用暴力做法做一次，题目限定只能使用 SET 或者 java.util.TreeSet，这个就比较简单了，只需要注意一下 corner case，然后注意参数不合法的时候抛异常。

其的搜索和插入的算法与 BST 的算法相似，但是在根结点处，我们使用 x 坐标来判断大小，如果要插入的点的 x 坐标比在根结点的点小，向左移动，否则向右移动；然后在下一个级别，我们使用 y 坐标来判断大小，如果要插入的点的 y 坐标比结点中的点小，则向左移动，否则向右移动；然后在下一级，继续使用 x 坐标，依此类推……

由此，我们可以得到下图：

Kd-Trees 插入示意

相对于 BST 的主要优势在于，它支持范围搜索和最近邻居搜索的高效实现。每个节点对应于单位正方形中与轴对齐的矩形，该矩形将其子树中的所有点都包含在内。根结点对应整个单位正方形，根的左、右子元素对应于两个矩形，该两个矩形被根结点的 x 坐标分开，以此类推……

由此，我们可以得到范围搜索和最近邻居搜索的思想思路。

进行范围搜索时，从根结点开始，递归地搜索左右子树，若查询矩形不与该结点对应的矩形相交，那么就不需要探索该节点及其子树。子树只有在可能包含查询矩形中包含的点时才被搜索。

进行最近邻居搜索时，从根结点开始，递归地搜索左右子树，如果到目前为止发现的最近点比查询点与结点对应的矩形之间的距离更近，则不需要探索该结点及其子树。也就是说，仅当一个结点可能包含一个比目前发现的最佳结点更接近的点时，才进行搜索。

这样的剪枝规则，依赖于能否快速找到附近的点。因此，我们需要注意在递归代码中，当有两个可能的子树的时候，总是选择位于分隔线同一侧的子树作为要探索的第一棵子树的查询点。这是因为在探索第一棵子树时发现的有可能是最近的点，将有利于探索第二棵子树时剪枝。

这里在实现的时候，递归先左先右当然都可以得到正确的结果，但是这里必须调整递归的顺序，才能达到剪枝的效果。

这是因为，如果左孩子包含 p，由于矩形是越来越小的，所以若点在某个 node 的矩形内被包含，则该 node 的 p 离这个所求 p 的距离就可能越小。min 越小，那么剪枝的效果就越明显，因为越来越多的就不需要再计算了。于是，应该始终优先去递归那个 contains(p) 的方向（因为有且只有可能要么是 left 要么还是 right）包含 p。

如果不进行剪枝，那么就算你的代码功底非常好，在规定时间内求得了正确解，没有超时，也一样不能通过测评：

 - student sequence of kd-tree nodes involved in calls to Point2D methods: 
 A D F I G B C E J - reference sequence of kd-tree nodes involved in calls to Point2D methods: 
 A D F I G B C
 - failed on trial 1 of 1000

具体剪枝的策略就是，如果左孩子包含了目标点，那么就去左孩子，如果右孩子包含了目标点，那么就去右孩子。有可能左右孩子都不包含目标点，那么离谁近就去谁那。

// 先左先右当然都可以得到正确的结果，但是
// 这里必须调整递归的顺序，才能达到剪枝的效果
if (node.left != null && node.left.rect.contains(p)) {    
    // 如果左孩子包含 p，由于矩形是越来越小的，所以若点在某个 node 的矩形内被包含，则该 node 的 p 离这个所求 p 的距离就可能越小    
    // min 越小，那么剪枝的效果就越明显，因为越来越多的就不需要再计算了    
    // 于是，应该始终优先去递归那个 contains(p) 的方向（因为有且只有可能要么是 left 要么还是 right）包含 p    
    findNearest(p, node.left);    
    findNearest(p, node.right);
} else if (node.right != null && node.right.rect.contains(p)) {    
    // 如果右孩子包含就先去右边    
    findNearest(p, node.right);    
    findNearest(p, node.left);
} else {    
    // 也可能出现两个都不包含的情况，那么离谁近就先去谁那    
    // 注意调用时 null 的问题要特别处理，可以设置为无穷大    
    double toLeft = node.left != null ? node.left.rect.distanceSquaredTo(p) : Double.POSITIVE_INFINITY;    
    double toRight = node.right != null ? node.right.rect.distanceSquaredTo(p) : Double.POSITIVE_INFINITY;    
    if (toLeft < toRight) {        
        findNearest(p, node.left);        
        findNearest(p, node.right);    
    } else {        
        findNearest(p, node.right);        
        findNearest(p, node.left);    
    }
}

为了代码实现的方便，二叉树当然要用递归的写法啦。

课程提供了若干可视化工具用于调试。

draw() 函数的正确性将会大幅度提高 debug 的效率，所以这个函数一定要写的正确。

在可视化过程中，使用暴力法求解的答案会标注为红色，使用 KDTree 方法求解的会标注为蓝色。由于我们非常有信心，暴力法肯定是对的，所以可以用这个方法来检验 KdTree 的搜索是不是正确。