算法（1）：KD树的原理与基本代码

目录

前言

一、什么是KD树

二、为什么要用KD树

三、KD树的基本思路

四、KD树的几种情况分析

4.1 另一子空间不存在更近的点

4.2 另一子空间存在更近的点

4.3 小结

五、KD树的代码（二维点，python版本）

六、KD树的代码（多维版本）

6.1 python版本

七、KD树的应用

7.1 找目标平面或者空间中离目标点的最近点

7.2 找目标平面或者空间中离目标点的若干个最近点

八、参考资料

前言

由于在做项目时遇到平面内最近点求解的问题，需要用到KD树简化算法，减少计算资源，因此本文记录这个过程中学到的知识。

一、什么是KD树

kd-tree（k-dimensional树的简称），kd树就是一种对k维空间中的实例点进行存储以便对其进行快速检索的树形数据结构，可以运用在k近邻法中，实现快速k近邻搜索。构造kd树相当于不断地用垂直于坐标轴的超平面将k维空间切分。k-d树是每个节点都为k维点二叉树。所有非叶子节点可以视作用超平面把空间分割成两半空间。节点左边的子树代表在超平面左边的点，节点右边的子树代表在超平面右边的点。

二、为什么要用KD树

主要是为了快速计算，节省计算资源。如前言中提到的例子，计算平面中一堆点中哪一个和目标点最近的问题，如果使用常规方法需要计算所有点与目标点的距离，然后进行比较，这样在平面中点非常多的时候非常浪费计算资源，速度也很慢。为了解决这个问题，可以在这个过程中引入KD树进行算法上的简化。

三、KD树的基本思路

下面是一个学习KD树的经典案例：

在二维平面上有以下六个点：(2,3），（5,4），（9,6），（4,7），（8,1），（7,2）

KD树的目的要确定图1中这些分割空间的分割线，确定步骤如下：

（1）第一步，将上述六个点的点集按照二维平面的第一维（x轴）对数据进行排序，排序结果为（2,3），（4,7），（5,4），（7,2），（8,1），（9,6）

（2） 第二步，取得上述点集的中位数处的点，偶数个数的中位数一般取大的那个，在上述点集中取到的为（7,2），在该点处对平面进行划分，如图1中（7,2）位置画上一条竖线

（3）第三步，将由（7,2）划分的左右两平面中的剩余点按照二维平面的第二维（y轴）对数据排序，排序结果为左枝：（2,3），（5,4），（4,7）和右枝（8,1），（9,6）

（4）第四步，将左枝和右枝分别取中位数点作为新的节点，按照第一维继续排序，直致每一个子枝上只剩一个点，这个点被称作叶子。

上述过程画成图如下：

需要注意的是，每次划分之后排序的逻辑是按照纬度一次类推的，在二维平面中就是按x,y,x,y,x,y....轴排序，如果在三维空间中就是x,y,z,x,y,z,x,y,z....轴排序。

进行完上述操作后就将数据划分完毕，可以将目标点放到该KD树中比较了。

四、KD树的几种情况分析

4.1 另一子空间不存在更近的点

案例如下：

初始点集仍为(2,3），（5,4），（9,6），（4,7），（8,1），（7,2）六个点

目标点为（2.1,3.1）

依据上述KD树的搜索逻辑，操作如下：

（1）将初始点击根据第一维（X坐标）排序为（2,3）（4,7）（5,4）（7,2）（8,1）（9，6），其中位数取4，也就是（7,2）点，在上图中 X=7 处画上分割线

（2）判断目标点的X坐标2.1小于（7,2），因此取划分后的左侧空间中的点（2，3）（4,7）（5,4）

（3）将上述点集根据第二维（Y坐标）排序为（2,3）（5,4）（4,7），其中位数是2，即（5,4）点，在平面中画上该线

（4）判断目标点的Y坐标3.1小于（5,4），因此在划分后的下方空间中

（5）此时改空间中仅有（2,3）一个点，到此结束划分空间的操作

（6）计算此时得到的最后点与目标点之间的距离为，暂时记录下这个值tempDis

（7）下面开始进行回溯

计算目标点（2.1,3.1）与最后的节点（5,4）之间Y轴的距离差记为dist，计算得0.9，这个值大于≈0.1414，不可能在节点（5,4）的上侧空间有更近的点，因此在节点（5,4）下方空间中找到的点（2,3）已经是最近的点。继续回溯上一个节点（5,4），此时在X轴上的距离2.9>0.1414，不存在更近点的可能，因此回溯结束，最近的点为（2,3），最近的距离为。

实际情况中回溯时会碰到上述dist小于tempDis，有可能在上一节点的另一空间中存在更近点的可能，这种情况参见下面的案例。

4.2 另一子空间存在更近的点

案例如下：

初始点集为（5,3），（2.5,5），（8,4.5），（2,2），（3.5,8），（8,2.5），（5.5,7.5）

目标点为（4.5,7.5）

依据KD树的搜索逻辑，操作如下：

（1）将初始点击根据第一维（X坐标）排序为（2,2）（2.5,5）（3.5,8）（5,3）（5.5,7.5）（8，4.5），（8,2.5），其中位数取4，也就是（5,3）点，在上图中 X=5 处画上分割线

（2）判断目标点的X坐标4.5小于（5,3），因此取划分后的左侧空间中的点（2，2）（2.5,5）（3.5,8）

（3）将上述点集根据第二维（Y坐标）排序为（2，2）（2.5,5）（3.5,8），其中位数是2，即（2.5,5）点，在平面中画上该线Y=5

（4）判断目标点的Y坐标7.5大于（2.5,5），因此在划分后的上方空间中

（5）此时改空间中仅有（2.5,8）一个点（点D），到此结束划分空间的操作

（6）计算此时得到的最后点与目标点之间的距离为≈2.06，暂时记录下这个值tempDis

（7）下面开始进行回溯

计算目标点（4.5,7.5）与最后的节点B（2.5,5）之间Y轴的距离差记为dist，计算得2.5，这个值大于2.06，不可能存在节点（2.5,5）的下侧空间有更近的点；继续回溯上一个节点A（5,3），此时在X轴上的距离dist（5-4.5=0.5）<2.06，所以在A的右侧空间中有可能存在更近的点，我们进入到节点A的右侧空间，继续之前的操作，发现存在点E（5.5,7.5），其与目标点S之间的距离为1,1<2.06，因此最近距离tempDis更新为1，最近点更新为（5.5,7.5），回溯结束。

4.3 小结

总结一下，在写KD树的代码之前需要有一些准备：

（1）定义一个节点的类，需要有该节点的点坐标、左枝、右枝、划分的维度（用于确认需要根据哪个维度来对坐标排序）

（2）确定递归的出口（当划分的点集中只有一个点时结束）

（3）计算中位数点的编号

（4）计算当前划分的维度

（5）计算两个点的欧式距离（勾股定理）

（6）回溯获得最近点的坐标

五、KD树的代码（二维点，python版本）

import math

pts = [(5,3),(2.5,5),(8,4.5),(2,2),(3.5,8),(8,2.5),(5.5,7.5)]  #点集
targetPt = (4.5,7.5)  #目标点

class Node():
    def __init__(self,pt,leftBranch,rightBranch,dimension):
        self.pt = pt
        self.leftBranch = leftBranch
        self.rightBranch = rightBranch
        self.dimension = dimension

class KDTree():
    def __init__(self,data):
        self.nearestPt = None
        self.nearestDis = math.inf
    
    def createKDTree(self,currPts,dimension):
        if(len(currPts) == 0):
            return None
        mid = self.calMedium(currPts)
        sortedData = sorted(currPts,key=lambda x:x[dimension])
        leftBranch = self.createKDTree(sortedData[:mid],self.calDimension(dimension))
        rightBranch = self.createKDTree(sortedData[mid+1:],self.calDimension(dimension))
        return Node(sortedData[mid],leftBranch,rightBranch,dimension)

    def calMedium(self,currPts):
        return len(currPts) // 2

    def calDimension(self,dimension):
        return (dimension+1)%2

    def calDistance(self,p0,p1):
        return math.sqrt((p0[0]-p1[0])**2+(p0[1]-p1[1])**2)

    def getNearestPt(self,root,targetPt):
        self.search(root,targetPt)
        return self.nearestPt,self.nearestDis
        
    def search(self,node,targetPt):
        if node == None:
            return
        dist = node.pt[node.dimension] - targetPt[node.dimension]
        if(dist > 0):#目标点在节点的左侧或上侧
            self.search(node.leftBranch,targetPt)
        else:
            self.search(node.rightBranch,targetPt)
        tempDis = self.calDistance(node.pt,targetPt)
        if(tempDis < self.nearestDis):
            self.nearestDis = tempDis
            self.nearestPt = node.pt
        #回溯
        if(self.nearestDis > abs(dist)):
            if(dist > 0):
                self.search(node.rightBranch,targetPt)
            else:
                self.search(node.leftBranch,targetPt)

if __name__ == "__main__":
    kdtree = KDTree(pts) 
    root = kdtree.createKDTree(pts,0)   
    pt,minDis = kdtree.getNearestPt(root,targetPt)
    print("最近的点是",pt,"最小距离是",str(minDis))

六、KD树的代码（多维版本）

6.1 python版本

import math

pts = []  #点集，任意维度的点集
targetPt =   #目标点，任意维度的点

class Node():
    def __init__(self,pt,leftBranch,rightBranch,dimension):
        self.pt = pt
        self.leftBranch = leftBranch
        self.rightBranch = rightBranch
        self.dimension = dimension

class KDTree():
    def __init__(self,data):
        self.nearestPt = None
        self.nearestDis = math.inf
    
    def createKDTree(self,currPts,dimension):
        if(len(currPts) == 0):
            return None
        mid = self.calMedium(currPts)
        sortedData = sorted(currPts,key=lambda x:x[dimension])
        leftBranch = self.createKDTree(sortedData[:mid],self.calDimension(dimension))
        rightBranch = self.createKDTree(sortedData[mid+1:],self.calDimension(dimension))
        return Node(sortedData[mid],leftBranch,rightBranch,dimension)

    def calMedium(self,currPts):
        return len(currPts) // 2

    def calDimension(self,dimension): # 区别就在于这里，几维就取余几
        return (dimension+1)%len(targetPt)

    def calDistance(self,p0,p1):
        return math.sqrt((p0[0]-p1[0])**2+(p0[1]-p1[1])**2)

    def getNearestPt(self,root,targetPt):
        self.search(root,targetPt)
        return self.nearestPt,self.nearestDis
        
    def search(self,node,targetPt):
        if node == None:
            return
        dist = node.pt[node.dimension] - targetPt[node.dimension]
        if(dist > 0):#目标点在节点的左侧或上侧
            self.search(node.leftBranch,targetPt)
        else:
            self.search(node.rightBranch,targetPt)
        tempDis = self.calDistance(node.pt,targetPt)
        if(tempDis < self.nearestDis):
            self.nearestDis = tempDis
            self.nearestPt = node.pt
        #回溯
        if(self.nearestDis > abs(dist)):
            if(dist > 0):
                self.search(node.rightBranch,targetPt)
            else:
                self.search(node.leftBranch,targetPt)

if __name__ == "__main__":
    kdtree = KDTree(pts) 
    root = kdtree.createKDTree(pts,0)   
    pt,minDis = kdtree.getNearestPt(root,targetPt)
    print("最近的点是",pt,"最小距离是",str(minDis))

七、KD树的应用

7.1 Unity找目标平面或者空间中离目标点的最近点

同上相似，不再赘述。

7.2 找目标平面或者空间中离目标点的若干个最近点

八、参考资料

[1] https://www.cnblogs.com/bambipai/p/8435797.html

[2] kd-tree_百度百科文章来源地址https://www.toymoban.com/news/detail-718433.html

到了这里，关于算法（1）：KD树的原理与基本代码的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！