KD-tree的原理以及构建与查询操作的python实现

原文地址http://blog.csdn.net/u010551621/article/details/44813299#comments

前几天小组讨论会上展示了kd-tree（k－dimension tree），感觉这玩意儿还挺有用的，所以学习了一下它的原理，然后把其中的构建kd-tree以及对应的查询操作实现了一下，现在跟大家分享一下

首先说一下什么是kd-tree把

不过首先得说一下bst(二叉查找树)，递归定义如下：如果左子树上的节点存储的数值都小于根节点中存储的数值，并且右子树上的节点存储的数值都大于根节点中存储的数值，那么这样的二叉树就是一颗二叉查找树

有了bst的概念，那么kd-tree就 容易理解多了，首先kd-tree的节点中存储的数值是一个k维的数据点，而bst的节点中存储的可以视为是1维的数据点，kd-tree与bst不同的地方在于进行分支决策的时候，还需要选择一个维度的值进行比较，选择哪个维度呢？每个节点还需要维护一个split变量，表示进行分支决策的时候，选择哪个维度的值进行比较，现在给出一个kd-tree节点的定义

[python] view
plain copy

class KD_node:  

    def __init__(self, point=None, split=None, LL = None, RR = None):  

        """ 

        point:数据点 

        split:划分域 

        LL, RR:节点的左儿子跟右儿子 

        """  

        self.point = point  

        self.split = split  

        self.left = LL  

        self.right = RR  

point就代表节点存储的k维数据点，left,right分别代表指向左右儿子的指针，split代表划分维度，在节点进行划分之前，我们需要确定划分维度，那么怎么选择划分维度呢，这又要从kd-tree的用途开始说起了

kd-tree是一种对高维空间的数据点进行划分的特殊数据结构，主要应用就是高维空间的数据查找，如：范围搜索和K近邻(knn)搜索，范围搜索就是给定查询点和距离阈值，获取在阈值范围内的所有数据点；knn搜索就是给定查询点和搜索点的数目n，查找出到搜索点最近的n个点的数目；

以上这两种搜索如果通过传统方法来实现，那么最坏情况下可能会穷举 数据急中的所有点，这种方法的缺点就是完全没有利用到数据集中蕴藏的结构信息，当数据点很多时，搜索效率不高；

事实上，实际数据集中的点一般时呈簇状分布的，所以，很多点我们是完全没有必要遍历的，索引树的方法就是对将要搜索的点进行空间划分，空间划分可能会有重叠，也可能没有重叠，kd-tree就是划分空间没有重叠的索引树

这样说可能有一点乱，那我还是以“二分查找”作为引入吧

如果给你一组数据   9 1 4 7 2 5 0 3 8

让你查找8，如果你挨个查找，那么将会把数据集都遍历一遍，

如果你排一下序那现在数据集就变成了：0 1 2 3 4 5 6 7 8 9，其实我们进行了很多没有必要的查找，

如果我以5为分界点，那么数据点就被分为了 两个“簇” （0 1 2 3 4）和（6 7 8 9），如果我要查找8，我根本久没有必要进入第一个簇，直接进入第二个簇进行查找，经过2次操作之后，就可以找到8了

把二分查找中的数据点换成k维数据点，这样的划分就变成了我们刚才说的空间划分，所以在这里要搞清楚，空间划分就是把数据点分类，“挨得近”的数据点就在一个空间里面

好 现在回到刚才的划分维度的选择上，因为我要尽可能将相似的点放在一颗子树里面，所以kd-tree采取的思想就是计算所有数据点在每个维度上的数值的方差

然后方差最大的维度就作为当前节点的划分维度，这样做的原理其实就是：方差越大，说明这个维度上的数据波动越大，也就说明了他们就越不可能属于同一个空间，需要在这个维度上对点进行划分，这就是kd-tree节点选择划分维度的原理
先贴一张kd-tree的图



途中每个节点代表划分点，标示维黑体的维度就是节点的划分维度，可以看到对于任意节点来说，如果给定划分维度split, 它的左子树上的节点在split维度上的值一定比它在split维度上的值要小，右子树上的节点在split维度上的值一定相应要大，所以说kd-tree实际上就是bst在多维空间上的拓展

好，扯了那么多废话，举个例子来说一下kd-tree的构造

现在假设我有若干个二维空间的数据点（横向为x轴，纵向为y轴）



通过第一次方差的计算，我们发现x维度上的方差比较大，所以，我们先选x轴为划分维度，得到下面的点，黄色的点代表分割点，这里要说明一下，分割点（也就是节点存储的数据节点）一般取在分割维度上的值为中间值的点，下图就是选了x维度上的值为中值的点作为切割点



现在我们又对x<x0 和x>=x0空间进行划分，先看x>=x0这个子空间，很明显,y轴方向上的数据波动要比x轴方向上的数据波动更大，所以这个空间中我们选的划分维度为y维度,红色节点为分割点



我们按照上面的方法，持续对空间中的点进行划分，直到每个空间中只有一个点，这样，一棵kd-tree就构成了



根据上面的介绍，黄色的节点就代表kd-tree的根节点，也就是第一个分割点；红色的点代表位于第二层上的节点，剩下的以此类推

好了，现在附上创建kd-tree的python代码

[python] view
plain copy

def createKDTree(root, data_list):  

    """ 

    root:当前树的根节点 

    data_list:数据点的集合(无序) 

    return:构造的KDTree的树根 

    """  

    LEN = len(data_list)  

    if LEN == 0:  

        return  

    #数据点的维度  

    dimension = len(data_list[0])  

    #方差  

    max_var = 0  

    #最后选择的划分域  

    split = 0;  

    for i in range(dimension):  

        ll = []  

        for t in data_list:  

            ll.append(t[i])  

        var = computeVariance(ll)  

        if var > max_var:  

            max_var = var  

            split = i  

    #根据划分域的数据对数据点进行排序  

    data_list.sort(key=lambda x: x[split])  

    #选择下标为len / 2的点作为分割点  

    point = data_list[LEN / 2]  

    root = KD_node(point, split)  

    root.left = createKDTree(root.left, data_list[0:(LEN / 2)])  

    root.right = createKDTree(root.right, data_list[(LEN / 2 + 1):LEN])  

    return root  

  

  

def computeVariance(arrayList):  

    """ 

    arrayList:存放的数据点 

    return:返回数据点的方差 

    """  

    for ele in arrayList:  

        ele = float(ele)  

    LEN = len(arrayList)  

    array = numpy.array(arrayList)  

    sum1 = array.sum()  

    array2 = array * array  

    sum2 = array2.sum()  

    mean = sum1 / LEN  

    #D[X] = E[x^2] - (E[x])^2  

    variance = sum2 / LEN - mean**2  

    return variance  

说完了kd-tree的构建，现在再来说一下如何利用kd-tree进行最近邻的查找

基本的查找思路是这样的：

1.二叉查找：从根节点开始进行查找，直到叶子节点；在这个过程中，记录最短的距离，和对应的数据点；同时维护一个栈，用来存储经过的节点

2.回溯查找：通过计算查找点到分割平面的距离(这个距离比较的是分割维度上的值的差，并不是分割节点到分割平面上的距离，虽然两者的值是相等的)与当前最短距离进行比较，决定是否需要进入节点的相邻子空间进行查找，为什么需要这个判断呢，我举一个例子就大家可能就能明白了



途中的黑点为kd-tree中的数据点，五角星为查询点，我们通过kd-tree的分支决策会将它分到坐上角的那部分空间，但并不是意味着它到那个空间中的点的距离最近

我们首先扫描到叶子节点，扫描的过程中记录的最近点为p（5，4），最短距离为d, 现在开始回溯，假设分割的维度为ss，其实回溯的过程就是确定是否有必要进入相邻子空间进行搜索，确定的依据就是当前点到最近点的距离d是否大于当前点到分割面（在二维空间中实际上就是一条线）的距离L，如果d < L,那么说明完全没有必要进入到另一个子空间进行搜索，直接继续向上一层回溯；如果有d > L,那么说明相邻子空间中可能有距查询点更近的点

python实现的代码如下：

[python] view
plain copy

def findNN(root, query):  

    """ 

    root:KDTree的树根 

    query:查询点 

    return:返回距离data最近的点NN，同时返回最短距离min_dist 

    """  

    #初始化为root的节点  

    NN = root.point  

    min_dist = computeDist(query, NN)  

    nodeList = []  

    temp_root = root  

    ##二分查找建立路径  

    while temp_root:  

        nodeList.append(temp_root)  

        dd = computeDist(query, temp_root.point)  

        if min_dist > dd:  

            NN = temp_root.point  

            min_dist = dd  

        #当前节点的划分域  

        ss = temp_root.split  

        if query[ss] <= temp_root.point[ss]:  

            temp_root = temp_root.left  

        else:  

            temp_root = temp_root.right  

    ##回溯查找  

    while nodeList:  

        #使用list模拟栈，后进先出  

        back_point = nodeList.pop()  

        ss = back_point.split  

        print "back.point = ", back_point.point  

        ##判断是否需要进入父亲节点的子空间进行搜索  

        if abs(query[ss] - back_point.point[ss]) < min_dist:  

            if query[ss] <= back_point.point[ss]:  

                temp_root = back_point.right  

            else:  

                temp_root = back_point.left  

  

            if temp_root:  

                nodeList.append(temp_root)  

                curDist = computeDist(query, temp_root.point)  

                if min_dist > curDist:  

                    min_dist = curDist  

                    NN = temp_root.point  

    return NN, min_dist  

  

  

def computeDist(pt1, pt2):  

    """ 

    计算两个数据点的距离 

    return:pt1和pt2之间的距离 

    """  

    sum = 0.0  

    for i in range(len(pt1)):  

        sum = sum + (pt1[i] - pt2[i]) * (pt1[i] - pt2[i])  

    return math.sqrt(sum)  

为了验证创建的树是否正确以及最后的距离度量是否正确，我分别使用了树的前序遍历和knn来对比运行的结果

[python] view
plain copy

def preorder(root):  

    """ 

    KDTree的前序遍历 

    """  

    print root.point  

    if root.left:  

        preorder(root.left)  

    if root.right:  

        preorder(root.right)  

  

  

def KNN(list, query):  

    min_dist = 9999.0  

    NN = list[0]  

    for pt in list:  

        dist = computeDist(query, pt)  

        if dist < min_dist:  

            NN = pt  

            min_dist = dist  

    return NN, min_dist