kd树的构造和搜索(超详细)

kd树的构造和搜索(举例+手写代码)

最近在学机器学习和NLP，学到KNN时遇到了一个很难的数据结构-kd树。用来对高维数据进行存储和搜索。KNN概述见上篇文章。
本文将对kd树的构造和搜索做非常详细的介绍。其中构造还是很简单的，难在搜索上，尤其是怎么对父节点进行回溯。
kd（k-dimensional）树解决的是为k维数据建立索引的问题。
已知样本空间和一个目标向量如何快速搜索到目标向量的最近邻？
常规方法是线性表+线性搜索。复杂度o(n)，数据量很大时很耗时。
而kd树理论上讲是o(log(n))。
通过空间换时间的方法减小时间复杂度，来对样本数据建立索引。
类似于平衡二叉树，不过是对k维数据进行划分。

Kd树的构造

假设特征空间T:

构造kd树流程：
输入：n个对象的特征向量组成的特征空间T={$x_0,x_1,\dots ,x_N$}
输出：n个对象的节点组成的kd树

1. 初始化：以所有训练集作为样本空间，l=0；
2. 构造根节点：选取样本空间内$x^l$的中位数a，对应样本为$x_i$，将$x_i$置为根节点，左子节点的样本空间为特征$x^l$小于a的子空间，右子节点的样本空间为特征$x^l$大于a的子空间。
3. 递归：以左右子节点作为根节点，$l=(l+1)\%n$ ;重复步骤2。
4. 递归终点：重复直到根节点样本空间为空，即所有样本被划分完毕。

步骤3中$l=(l+1)\%n$：+1是因为要对下一个特征进行划分，取模是因为对于特征少而样本多的数据而言，划分完一轮$(x^0,x^n)$后，可能还有样本没有划分完，由此我们可以再从特征$x^0$开始，直到样本均被划分。
此外，$depth\%n=l$，该结论在深度搜索kd树时有重要作用

举例：
假设有如下特征空间（3个特征，5个样本）：

ka树构造算法（代码思路）：

1. 初始化样本空间： T = { x 0 , x 1 , x 2 , x 3 , x 4 } T=\{x_0,x_1,x_2,x_3,x_4\} T={x0​,x1​,x2​,x3​,x4​}，$l=0$.
2. 对于样本空间T， x ( l = 0 ) = { 1 , 2 , 4 , 9 , 4 } x^{(l=0)}=\{1,2,4,9,4\} x(l=0)={1,2,4,9,4},中位数为4，对应样本为$x_4$，由此根节点为$x_4$，左子节点样本空间为 T 1 = { x 0 , x 1 } T_1=\{x_0,x_1\} T1​={x0​,x1​}，右子节点样本空间为 T 2 = { x 2 , x 3 } T_2=\{x_2,x_3\} T2​={x2​,x3​}，$l+=1$.
3. 对于样本空间$T_1$， x ( l = 1 ) = { 6 , 9 } x^{(l=1)}=\{6,9\} x(l=1)={6,9}，中位数有两个，那么我们选取其中一个，比如6（计算机特性，n/2一般都是前者，懂得都懂），对应样本为$x_0$，由此左子节点样本空间为 T 3 = { } T_3=\{\} T3​={}，右子节点样本空间为 T 4 = { x 1 } T_4=\{x_1\} T4​={x1​}；$l+=1$;
   3.1. 对于$T_3$={}，return；对于$T_4$= { x 1 } \{x_1\} {x1​}， $x^{(l=2)}$={2}，中位数为2，对应样本为$x_1$，左右子节点样本空间均为空，return。
4. 对于样本空间$T_2$， x ( l = 1 ) = { 1 , 3 } x^{(l=1)}=\{1,3\} x(l=1)={1,3}，同理中位数为1，对应样本为$x_2$，由此左子节点样本空间为 T 5 = { } T_5=\{\} T5​={}，右子节点样本空间为 T 6 = { x 3 } T_6=\{x_3\} T6​={x3​}。$l+=1$;
   4.1. 对于 T 5 = { } T_5=\{\} T5​={}，return；对于 T 6 = { x 3 } T_6=\{x_3\} T6​={x3​}， x ( l = 2 ) = { 5 } x^{(l=2)}=\{5\} x(l=2)={5}，中位数为2，对应样本为$x_3$，左右子节点样本空间均为空，return。

kd树如下：

kd树的递归构造函数（文末有完整代码）：

def build(arr,l,father):
    if len(arr)==0: #样本空间为空则返回
        return ;

    #找x^l的中位数和对应特征向量,即arr[:][l]的中位数及arr[x][:]
    size = len(arr)
    # 直接对arr进行排序，因为要得到特征向量和划分子空间，由此直接对arr排序最便捷
    # 对二维数组排序，排序值为l列（升序）：
    arr.sort(key=(lambda x:x[l]))

    mid = int((size-1)/2)

    root = TreeNode(arr[mid],l)
    root.father = father
    root.left=build(arr[0:mid],(l+1)%n,root)  #0:mid不包括mid，即[0,mid)
    root.right=build(arr[mid+1:],(l+1)%n,root)
    # print(root,root.father)
    return root

kd树的搜索

输入：kd树，目标对象x的特征向量
输出：x的k-近邻

先了解DFS算法（明确一点，该算法只是找个叶子结点而已，并不是找最近邻的算法）：

我们找到的这个“当前最近点”是否就是全局距离最小的点？
当然不一定，举个简单的例子：

如上图，假设目标节点target=(3,1,4)，根据dfs，当前最近点为(1,6,2)，但是父节点(4,2,6)和兄弟(5,1,4)明显更近。如果父节点(4,2,6)上面还有父节点和兄弟节点呢，是不是也有可能比(1,6,2)更近？
那怎么在全局找到最近邻呢？利用贪心和回溯思想!

为了方便讨论，有如下定义和方法（认真看完！）：

1. 当前最近点：当前离target的距离最近的点称为当前最近点 最近距离：当前最近点与target的距离称为最近距离
2. 超球体：以target为球心，最近距离为半径的球称为超球体
3. 结点的超平面：在划分点上，平行划分维度的平面。比如上图，我们第一次划分的时候以第一维的4为划分点，那么$x_4$的超平面就是x=4的平面
4. 判断超球体和超平面相交的方法：
   |超平面上划分点的对应维度的值 – 球心在对应维度的值| < R（最近距离）
   如上面的例子，超球体的球心：(1,6,2) ，半径：5.74，$x_4$的超平面：x=4，划分点：(4,2,6)，对应维度：1。
   带入上述公式：|4-1| =3 < 5.74 因此 $x_4$的超平面与超球体相交。

因此，我们通过dfs算法找到一个当前最近点之后，还需要讨论父节点和兄弟节点是否比当前最近点离target更近。并利用递归方法，回溯直到根节点。
（详细算法描述后文再说）

• 等等？这样的话不就需要判断所有结点了吗，和顺序搜索有什么区别？
• 错了，kd树之所以能够高效搜索，是因为如果父节点的超平面不和超球体相交的话，就不需要再讨论兄弟结点了!
• 通俗地理解就是：你有若干哥哥和若干弟弟（假设身高和年龄正相关），其中一个弟弟想和你的一个哥哥比身高，但是弟弟还没你高，也就没必要叫哥哥来比了。（转换为高维数据就是kd树的构造思想）

kd树的查找算法的完整描述：

1. DFS算法：从根节点出发，递归找到一个叶子结点。
   递归方法：若target当前维度值小于当前结点对应维度值，递归左子树；否则右子树。
   深度和维度的关系：$depth\%n=l$ ，其中n为特征的维度(个数)， depth为当前结点的深度（根的深度为1），$l$为当前深度的维度
2. 设置一个STACK，用来保存DFS的路径，以便回溯。
3. 先利用DFS找到一个当前最近点，无论这个点离真正的最近邻有多远。
4. 对STACK的元素进行回溯：
   4.1 保存好当前节点(栈顶元素)、父亲节点、兄弟节点
   4.2 出栈一个
   4.3 判断当前节点离target的距离是否小于最近距离：
   小于：更新当前最近点和最近距离
   Else ：4.4
   4.4 判断超球体和父节点的超平面是否相交：
   相交：则从兄弟节点开始DFS并把路径入栈。
   不相交：4.1
5. 回溯直到栈顶元素为根节点

4.3是为了利用STACK更新最近点和最近距离，4.4是为了更新STACK，理解了这一点就不难理解上述算法了！
还有一点：4.4中需要从兄弟节点开始DFS并将递归路径压入栈，这就会产生兄弟结点之间不断地入栈的问题，因此需要设置一个标志，表示如果当前结点进行过DFS，那么兄弟节点不再进行判断。

举例：

高维不好观察，所以举二维的例子：

我们利用DFS找到当前最近点是D点，STACK=[A,B,D]
对STACK进行回溯：

1. 当前节点:D，父节点:B, 兄弟节点:F，出栈一个，STACK=[A,B]
   我们先判断当前节点D点是否比D更接近S（这个是程序的局限性），再判断：B点的超平面（图中红线）不和超球体相交，那么我们就不需要再判断兄弟结点F了。回溯。
2. 当前节点:B，父节点:A, 兄弟节点:C，出栈一个，STACK=[A]
   当前节点B不比D更近，但是A的超平面和超球体相交了，因此我们有必要讨论兄弟节点C，从C点开始DFS，路径:C->E。STACK=[A,C,E] 。回溯。
3. 当前节点:E，父节点:C, 兄弟节点:G，出栈一个，STACK=[ A,C]
   当前节点E比D近，更新当前最近点和最近距离。C的超平面不与超球体相交，因此不必判断G。
4. 当前节点:C，父节点:A, 兄弟节点:B，出栈一个，STACK=[ A]
   当前节点C不比D近，且当前节点C递归过，因此不必再判断父节点和兄弟节点。
5. 当前节点:A，当前节点为根节点，递归结束！

最近点为E。

源码

强烈建议对照源码理解：

class TreeNode:
    def __init__(self, s , d):
        self.vec = s  # 特征向量
        self.Dimension = d  #即划分空间时的特征维度
        self.left = None  # 左子节点
        self.right = None  # 右子节点
        self.father = None # 父节点（搜索时需要往回退）

    def __str__(self):
        return str(self.vec)  # print 一个 Node 类时会打印其特征向量

def dis(arr1,arr2):
    res=0;
    for a,b in zip(arr1,arr2):
        res += (a-b)**2
    return res**0.5

def build(arr,l,father):
    if len(arr)==0: #样本空间为空则返回
        return ;

    #找x^l的中位数和对应特征向量,即arr[:][l]的中位数及arr[x][:]
    size = len(arr)
    # 直接对arr进行排序，因为要得到特征向量和划分子空间，由此直接对arr排序最便捷
    # 对二维数组排序，排序值为l列（升序）：
    arr.sort(key=(lambda x:x[l]))

    mid = int((size-1)/2)

    root = TreeNode(arr[mid],l)
    root.father = father
    root.left=build(arr[0:mid],(l+1)%n,root)  #0:mid不包括mid，即[0,mid)
    root.right=build(arr[mid+1:],(l+1)%n,root)
    # print(root,root.father)
    return root

def dfs(depth,root,father,stack):
    if root == None:
        return father;

    stack.append(root)

    if target[depth%n]<root.vec[depth%n]: #depth%n=l,l为特征的上标
        return dfs(depth + 1, root.left, root,stack)
    else :
        return dfs(depth+1,root.right,root,stack)
    return root;


Featureset=[[1,6,2],
            [2,9,3],
            [5,1,4],
            [9,4,7],
            [4,2,6],
            [6,3,5],
            [7,2,5],
            [9,1,4]]
n = len(Featureset[0][:]) #特征的维度(个数)
Features = [1,3,4,2,5,5,2,7]
temp=Featureset

#递归构造kd树
root =build(temp,0,None)

#深度搜索kd树，找到一个当前最近点
target = [1,4,5]
stack = []
nearest = dfs(0,root,root.father,stack) #找到最近邻
nearest_dis = dis(nearest.vec,target) #nearest_dis为当前最近邻离target的距离，即最小距离，也是超球体的半径

visited={} #用来判断兄弟节点是否已经讨论过


#利用stack进行回溯，而非递归
while stack[-1]!=root:
    print('STACK:',end=' ')
    for i in stack : print(i,end=' ')
    print()

    #先定义好当前节点cur、父亲节点father、兄弟节点bro
    cur = stack[-1]
    stack.pop();
    father = cur.father
    bro = father.left
    if father.left == cur:
        bro = father.right

    # 如果当前节点与target的距离小于最近距离，则更新最近结点和最近距离
    if dis(cur.vec,target)<nearest_dis:
        nearest=cur
        nearest_dis=dis(cur.vec,target)

    # print(visited.get(hash(bro)))
    # 当前节点没有递归过
    if visited.get(hash(cur))==None:
    # 若超球体和父节点的超平面相交，相交则父节点的另一侧，即兄弟节点所在划分域可能存在更近的节点
        if father.vec[father.Dimension]-target[father.Dimension] < nearest_dis:
            visited.update({hash(bro): 'yes'})
            dfs(father.Dimension,bro,father,stack)

print("最近距离：",nearest_dis)
print("最近邻：",nearest)

运行结果：

初学者，若有错误，欢迎指正！