机器学习常用的KD-Tree深度讲解

1、从线段树到KD树在讲KD树之前，我们先来了解一下线段树的概念。线段树在机器学习领域 当中不太常见，作为高性能维护的数据结构，经常出现在各种算法比赛当中。 线段树的本质是一棵维护一段区间的平衡二叉树。比如下图就是一个经典的线 段树：2 3X 1从下图当中我们不难看出来，这棵线段树维护的是个区间内的最大值。 比如树根是8,维护的是整个区间的最大值，每一个中间节点的值都是以它为 树根的子树中所有元素的最大值。通过线段树，我们可以在的时间内计算出某一个连续区间的最大值。比如 我们来看下图：b当我们要求被框起来的区间中的最大值，我们只需要找到能够覆盖这个区 间的中间节点就行。我们可以发现被红框框起来的两

2、个节点的子树刚好覆盖这 个区间，于是整个区间的最大值，就是这两个元素的最大值。这样，我们就把 一个需要查找的问题降低成了，不但如此，我们也可以做到复杂度内的更新， 也就是说我们不但可以快速查询，还可以更新线段当中的元素。当然线段树的 应用非常广泛，也有许多种变体，这里我们不过多深入，感兴趣的同学可以期 待一下周三的算法与数据结构专题，在之后的文章当中会为大家分享线段树的 相关内容。在这里，我们只需要有一个大概的印象，线段树究竟完成的是什么 样的事情即可。线段树维护的是一个线段，也就是区间内的元素，也就是说维护的是一个一维的序列。如果我们将数据的维度扩充一下，扩充到多维呢？是 的，你没有猜错，从

3、某种程度上来说，我们可以把KD-Tree看成是线段树拓展到多维空间当中的情况。KD-Tree 定义我们来看一下KD-Tree的具体定义，这里的 K指的是K维空间，D自然就 是dimension，也就是维度，也就是说 KD-Tree就是K维度树的意思。在我们 构建线段树的时候，其实是一个递归的建树过程，我们每次把当前的线段一分 为二，然后用分成两半的数据分别构建左右子树。我们可以简单写一下伪代码， 来更直观地感受一下：class Node:def _init_(self, value, lchild, rchild):self.value = valueself.lchild = lchilds

4、elf.rchild = rchild def build(arr):n = len( arr):left, right = arr: n /2, arrn /2:Ichild, rchild = self.build(left), self.build(right)return Node(max(lchild.value, rchild.value), Ichild, rchild)我们来看一个二维的例子，在一个二维的平面当中分布着若干个点。x轴。我们对我们首先选择一个维度将这些数据一分为二，比如我们选择 所有数据按照x轴的值排序，选出其中的中点进行一分为二。在这根线左右两侧的点被分成了两棵

5、子树，对于这两个部分的数据来说， 我们更换一个维度，也就是选择 y轴进行划分。一样，我们先排序，然后找到 中间的点，再次一分为二。我们可以得到：我们重复上述过程，一直将点分到不能分为止，为了能更好地看清楚，我 们对所有数据标上坐标（并不精确）。如果我们把空间看成是广义的区间，那么它和线段树的原理是一样的。最 后得到的也是一棵完美二叉树，因为我们每次都选择了数据集的中点进行划分， 可以保证从树根到叶子节点的长度不会超过叩N。我们代入上面的坐标之后, 我们最终得到的KD-Tree大概是下面这个样子：KD-Tree 建树在建树的过程当中，我们的树深每往下延伸一层，我们就会换一个维度作 为衡量标准。原

6、因也很简单，因为我们希望这棵树对于这K维空间都有很好的表达能力，方便我们根据不同的维度快速查询。在一些实现当中，我们会计算 每一个维度的方差，然后选择方差较大的维度进行切分。这样做自然是因为方 差较大的维度说明数据相对分散，切分之后可以把数据区分得更加明显。但我 个人觉得这样做意义不是很大，毕竟计算方差也是一笔开销。所以这里我们选 择了最朴素的方法一一轮流选择。也就是说我们从树根开始，选择第0维作为排序和切分数据的依据，然后到了树深为1的这一层，我们选择第一维，树深为2的这一层，我们选择第二维，以此类推。当树深超过了K的时候，我们就对树深取模。明确了这一点之后，我们就可以来写KD-Tree的建

7、树代码了，和上面二叉树的代码非常相似，只不过多了维度的处理而已。#外部暴露接口def build_model(self, dataset):self.root = self._build_model(dataset)#先忽略，容后再讲self.set_father(self.root, None)#内部实现的接口def _build_model(self, dataset, dep th=O): if len( dataset) = 0:return None#通过树深对K取模来获得当前对哪一维切分axis = depth % self.Km = len( dataset) / 2#根据axi

8、s这一维排序dataset = sorted(dataset, key=lambda x: xaxis)#将数据一分为二left, mid, right = dataset:m, datasetm, datasetm+1:#递归建树return KDTree.Node(midaxis,mid,axis,dep th,len( dataset),self._build_model(left, dep th+1), self._build_model(nght, dep th+1)这样我们就建好了树，但是在后序的查询当中我们需要访问节点的父节点， 所以我们需要为每一个节点都赋值指向父亲节点的指针。

9、这个值我们可以写在 建树的代码里，但是会稍稍复杂一些，所以我把它单独拆分了出来，作为一个 独立的函数来给每一个节点赋值。对于根节点来说，由于它没有父亲节点，所 以赋值为Nones我们来看下set_father当中的内容，其实很简单，就是一个 树的递归遍历：def set_father(self, no de, father):if node is None:return#赋值no de.father = father#递归左右self.set_father( no de.lchild, no de)self.set_father( no de.rchild, no de)快速批量查询KD-Tr

10、ee建树建好了肯定是要来用的，它最大的用处是可以在单次查询中 获得距离样本最近的若干个样本。在分散均匀的数据集当中，我们可以在的时 间内完成查询，但是对于特殊情况可能会长一些，但是也比我们通过朴素的方 法查询要快得多。我们很容易发现，KD-Tree 一个广泛的使用场景是用来优化KNN算法。我们在之前介绍KNN算法的文章当中曾经提到过，KNN算法在预测的时候需要遍 历整个数据集，然后计算数据集中每一个样本与当前样本的距离，选出最近的 K个来，这需要大量的开销。而使用KD-Tree，我们可以在一次查询当中直接查找到K个最近的样本，因此大大提升 KNN算法的效率。那么，这个查询操作又 是怎么实现的呢

11、？这个查询基于递归实现，因此对于递归不熟悉的小伙伴，可 能初看会比较困难，可以先阅读一下之前关于递归的文章。首先我们先通过递归查找到KD-Tree上的叶子节点，也就是找到样本所在 的子空间。这个查找应该非常容易，本质上来说我们就是将当前样本不停地与 分割线进行比较，看看是在分割线的左侧还是右侧。和二叉搜索树的元素查找 是一样的：def iter_dow n( self, no de, data):#如果是叶子节点，则返回if no de.lchild is None and no de.rchild is None:return node#如果左节点为空，则递归右节点if no de.lchi

12、ld is None:return self.iter_dow n(no de.rchild, data)#同理，递归左节点if no de.rchild is None:return self.iter_dow n(no de.rchild, data)#都不为空则和分割线判断是左还是右axis = no de.axisn ext_ node = no de.lchild if dataaxis = no de.bo un dray else no de.rchild return self.iter_dow n(n ext_ no de, data)我们找到了叶子节点，其实代表样本空间当中

13、的一小块空间。我们来实际 走一下整个流程，假设我们要查找3个点。首先，我们会创建一个候选集，用来存储答案。当我们找到叶子节点之后，这个区域当中只有一个点，我们把它 加入候选集。在上图当中紫色的x代表我们查找的样本，我们查找到的叶子节点之后， 在两种情况下我们会把当前点加入候选集。第一种情况是候选集还有空余，也 就是还没有满K个，这里的K是我们查询的数量，也就是3。第二种情况是当 前点到样本的距离小于候选集中最大的一个，那么我们需要更新候选集。这个 点被我们访问过之后，我们会打上标记，表示这个点已经访问过了。这个时候 我们需要判断，整棵树当中的搜索是否已经结束，如果当前节点已经是根节点 了，说明

14、我们的遍历结束了，那么返回候选集，否则说明还没有，我们需要继 续搜索。上图当中我们用绿色表示样本被放入了候选集当中，黄色表示已经访 问过。由于我们的搜索还没有结束，所以需要继续搜索。继续搜索需要判断样 本和当前分割线的距离来判断和分割线的另一侧有没有可能存在答案。由于叶 子节点没有另一侧，所以作罢，我们往上移动一个，跳转到它的父亲节点。此时候选集未满，我们加入候选集，标记 但是也没有另一侧的节点，所以也跳过。 我们执行同样的判断，发现此时候选集还我们计算距离并且查看候选集, 为已经访问过。它虽然存在分割线, 我们再往上，遍历到它的父亲节点, 有空余，于是将它继续加入答案：但是当我们判断到分割线

15、距离的时候，我们发现这一次样本到分割线的举 例要比之前候选集当中的最大距离要小，所以分割线的另一侧很有可能存在答 案：这里的d1是样本到分割线的距离，d2是样本到候选集当中最远点的距离。 由于到分割线更近，所以分割线的另一侧很有可能也存在答案，这个时候我们 需要搜索分割线另一侧的子树，一直搜索到叶子节点。我们找到了叶子节点，计算距离，发现此时候选集已经满了，并且它的距 离大于候选集当中任何一个答案，所以不能构成新的答案。于是我们只是标记 它已经访问过，并不会加入候选集。同样，我们继续往上遍历，到它的父节点:J比较之后发现，data到它的距离小于候选集当中最大的那个，于是我们更 新候选集，去掉距

16、离大于它的答案。然后我们重复上述的过程，直到根节点为 止。由于后面没有更近的点，所以候选集一直没有更新，最后上图当中的三个 打了绿标的点就是答案。我们把上面的流程整理一下，就得到了递归函数当中 的逻辑，我们用Python写出来其实已经和代码差不多了：def query (no de, data, an swers, K):#判断node是否已经访问过if no de.visited:#标记访问no de.visited = True#计算data到node中点的距离dis = cal_dis(data, node.point)#如果小于答案中最大值则更新答案if dis max(a nswer

17、s):an swers. up date( node.point)#计算data到分割线的距离dis = cal_dis(data, no de.s plit)#如果小于最长距离，说明另一侧还可能有答案if dis max(a nswers):#获取当前节点的兄弟节点brother = self.get_brother( no de)if brother is not None:#往下搜索到叶子节点，从叶子节点开始寻找leaf = self.iter_dow n( brother, data)if leaf is not None:retur n self.query(leaf, data,

18、an swers, K)#如果已经到了根节点了，退出if node is root:return an swers#递归父亲节点return seif.query (no de.father, data, an swers, K) else:if node is root:return an swersreturn seif.query (no de.father, data, an swers, K)最终写成的代码和上面这段并没有太多的差别，在得到距离之后和答案当 中的最大距离进行比较的地方，我们使用了优先队列。其他地方几乎都是一样 的，我也贴上来给大家感受一下：def _query_ ne

19、arest_k(self, no de, p ath, data, topK, K):#我们用set记录访问过的路径，而不是直接在节点上标记if node not in p ath:p ath.add( no de)#计算欧氏距离dis = KDTree.dista nce( no de.value, data)if (len (to pK) K or dis top K-1dista nee):top K.a ppen d( no de: no de, dista nee: dis)#使用优先队列获取topKtopK = hea pq.n smallest(K, topK, key=lamb

20、da x: xdista nee)axis = no de.axis#分割线都是直线，直接计算坐标差dis = abs(dataaxis - no de.bo un dray)if len( to pK) K or dis = top K-1dista nee:brother = self.get_brother( no de, p ath)if brother is not None:n ext_ node = self.iter_dow n( brother, data)if n ext_ node is not None:retur n self._query_ nearest_k (n

21、 ext_ no de, p ath, data, topK, K)if node = self.root:return topKretu rn self._query_ nearest_k (no de.father, p ath, data, topK, K) else:if node = self.root:return topKretu rn self._query_ nearest_k (no de.father, p ath, data, topK, K)这段逻辑大家应该都能看明白，但是有一个疑问是，我们为什么不在node里面加一个visited的字段，而是通过传入一个set来维护访问过的节点呢？这个逻辑只看代码是很难想清楚的，必须要亲手实验才会理解。如果在node当中加入一个字段当然也是可以的，如