KD-树(k-dimensional tree)是一种用于组织点在k维空间中的数据结构，主要用于各种搜索和优化任务，如最近邻搜索、范围搜索和k最近邻搜索。KD-树是二叉树的一种特殊形式，可以看作是二分搜索树(BST)的推广，但适用于多个维度。本文记录相关内容。

简介

kd树(k-dimensional树的简称)，是一种分割k维数据空间的数据结构，其中的每一个节点都是k维的数据，主要应用于多维空间关键数据的近邻查找(Nearest Neighbor)和近似最近邻查找(Approximate Nearest Neighbor)。

其实KDTree就是二叉查找树(Binary Search Tree，BST)的变种。二叉查找树的性质如下：

1. 若它的左子树不为空，则左子树上所有结点的值均小于它的根结点的值；
2. 若它的右子树不为空，则右子树上所有结点的值均大于它的根结点的值；
3. 它的左、右子树也分别为二叉排序树；

如果我们要处理的对象集合是一个K维空间中的数据集，我们首先需要确定是：怎样将一个K维数据划分到左子树或右子树？

在构造1维BST树类似，只不过对于Kd树，在当前节点的比较并不是通过对K维数据进行整体的比较，而是选择某一个维度d，然后比较两个K维数据在该维度 d上的大小关系，即每次选择一个维度d来对K维数据进行划分，相当于用一个垂直于该维度d的超平面将K维数据空间一分为二，平面一边的所有K维数据 在d维度上的值小于平面另一边的所有K维数据对应维度上的值。也就是说，我们每选择一个维度进行如上的划分，就会将K维数据空间划分为两个部分，如果我 们继续分别对这两个子K维空间进行如上的划分，又会得到新的子空间，对新的子空间又继续划分，重复以上过程直到每个子空间都不能再划分为止。以上就是构造 Kd-Tree的过程，上述过程中涉及到两个重要的问题：

1. 每次对子空间的划分时，怎样确定在哪个维度上进行划分；
2. 在某个维度上进行划分时，怎样确保建立的树尽量地平衡，树越平衡代表着分割得越平均，搜索的时间也就是越少。

1、在哪个维度上进行划分？

一种选取轴点的策略是 median of the most spread dimension pivoting strategy，统计样本在每个维度上的数据方差，挑选出对应方差最大值的那个维度。数据方差大说明沿该坐标轴方向上数据点分散的比较开。这个方向上，进行数据分割可以获得最好的平衡。

2、怎样确保建立的树尽量地平衡？

给定一个数组，怎样才能得到两个子数组，这两个数组包含的元素 个数差不多且其中一个子数组中的元素值都小于另一个子数组呢？方法很简单，找到数组中的中值(即中位数，median)，然后将数组中所有元素与中值进行 比较，就可以得到上述两个子数组。同样，在维度d上进行划分时，划分点(pivot)就选择该维度d上所有数据的中值，这样得到的两个子集合数据个数就基本相同了。

构建 Kd-Tree

1. 在K维数据集合中选择具有最大方差的维度k，然后在该维度上选择中值m为pivot对该数据集合进行划分，得到两个子集合；同时创建一个树结点node，用于存储；
2. 对两个子集合重复(1)步骤的过程，直至所有子集合都不能再划分为止；

示例

假设有 6 个二维数据点 ${(2,3)，(5,4)，(9,6)，(4,7)，(8,1)，(7,2)}$，数据点位于二维空间内(如下图中黑点所示)。

kd树算法就是要确定图中这些分割空间的分割线(多维空间即为分割平面，一般为超平面)。下面就要通过一步步展示kd树是如何确定这些分割线的。

1. 分别计算x，y方向上数据的方差，得知x方向上的方差最大；
2. 根据x轴方向的值2,5,9,4,8,7排序选出中值为7，所以该node中的data = (7,2)。这样，该节点的分割超平面就是通过(7,2)并垂直于x轴的直线x = 7；
3. 确定左子空间和右子空间。分割超平面x = 7将整个空间分为两部分，如下图所示。x < = 7的部分为左子空间，包含3个节点{(2,3)，(5,4)，(4,7)}；另一部分为右子空间，包含2个节点{(9,6)，(8,1)}。

4. 更新分割轴：一共就两个维度，所以，下一个维度是y轴。

5. 确定子节点：

   左节点：在左支中找到y轴的中位数 (5,4)，左支数据更新为{(2,3)}，右支数据更新为{(4,7)}

   右节点：在右支中找到y轴的中位数 (9,6)，左支数据更新为{(8,1)}，右支数据为null。

6. 更新分割轴：下一个维度为x轴。

7. 确定(5,4)的子节点：

   左节点：由于只有一个数据，所以，左节点为 (2,3)

   右节点：由于只有一个数据，所以，右节点为 (4,7)

8. 确定 (9,6) 的子节点：

   左节点：由于只有一个数据，所以，左节点为 (8,1)

   右节点：右节点为空。

kd-tree表示：

Kd-Tree 最近邻查找

在构建了完整的kd-tree之后，我们想要使用他来进行高维空间的检索。

1. 将查询数据Q从根结点开始，按照Q与各个结点的比较结果向下访问Kd-Tree，直至达到叶子结点。
   其中Q与结点的比较指的是将Q对应于结点中的k维度上的值与中值m进行比较，若Q(k) < m，则访问左子树，否则访问右子树。达到叶子结点时，计算Q与叶子结点上保存的数据之间的距离，记录下最小距离对应的数据点，记为当前最近邻点nearest和最小距离dis。

2. 进行回溯操作，该操作是为了找到离Q更近的“最近邻点”。即判断未被访问过的分支里是否还有离Q更近的点，它们之间的距离小于dis。
   如果Q与其父结点下的未被访问过的分支之间的距离小于dis，则认为该分支中存在离P更近的数据，进入该结点，进行(1)步骤一样的查找过程，如果找到更近的数据点，则更新为当前的最近邻点nearest，并更新dis。

如果Q与其父结点下的未被访问过的分支之间的距离大于dis，则说明该分支内不存在与Q更近的点。
回溯的判断过程是从下往上进行的，直到回溯到根结点时已经不存在与P更近的分支为止。

注：判断未被访问过的树分支中是否还有离Q更近的点，就是判断"Q与未被访问的树分支的距离|Q(k) - m|“是否小于"Q到当前的最近邻点nearest的距离dis”。从几何空间上来看，就是判断以Q为中心，以dis为半径超球面是否与未被访问的树分支代表的超矩形相交。

下面举两个例子来演示一下最近邻查找的过程。
假设我们的kd树就是上面通过样本集{(2,3)，(5,4)，(9,6)，(4,7)，(8,1)，(7,2)}创建的。

例一

查找 (2.1,3.1) 的最近邻。

如下图所示，红色的点即为要查找的点。通过二叉搜索，顺着搜索路径很快就能找到当前的最邻近点(2,3)。

在上述搜索过程中，产生的搜索路径节点有<(7,2)，(5,4)，(2,3)>。为了找到真正的最近邻，还需要进行’回溯’操作，首先以(2,3)作为当前最近邻点nearest，计算其到查询点Q(2.1,3.1)的距离dis为0.1414，然后回溯到其父节点(5,4)，并判断在该父节点的其他子节点空间中是否有距离查询点Q更近的数据点。以(2.1,3.1)为圆心，以0.1414为半径画圆，如图所示。发现该圆并不和超平面y = 4交割，即这里：|Q(k) - m|=|3.1 - 4|=0.9 > 0.1414，因此不用进入(5,4)节点右子空间中去搜索。

再回溯到(7,2)，以(2.1,3.1)为圆心，以0.1414为半径的圆更不会与x = 7超平面交割，因此不用进入(7,2)右子空间进行查找。至此，搜索路径中的节点已经全部回溯完，结束整个搜索，返回最近邻点(2,3)，最近距离为0.1414。

例二

查找点Q(2,4.5)

如下图所示，同样经过二叉搜索，可得当前的最邻近点(4,7)，产生的搜索路径节点有<(7,2)，(5,4)，(4,7)>。首先以(4,7)作为当前最近邻点nearest，计算其到查询点Q(2,4.5)的距离dis为3.202，然后回溯到其父节点(5,4)，并判断在该父节点的其他子节点空间中是否有距离查询点Q更近的数据点。以(2,4.5)为圆心，以为3.202为半径画圆，如图所示。发现该圆和超平面y = 4交割，即这里：|Q(k) - m|=|4.5 - 4|=0.5 < 3.202，因此进入(5,4)节点右子空间中去搜索。所以，将(2,3)加入到搜索路径中，现在搜索路径节点有<(7,2), (2, 3)>。同时，注意：点Q(2,4.5)与父节点(5,4)的距离也要考虑，由于这两点间的距离3.04 < 3.202，所以将(5,4)赋给nearest，并且dist=3.04。

接下来，回溯至(2,3)叶子节点，点Q(2,4.5)和(2,3)的距离为1.5，比距离(5,4)要近，所以最近邻点nearest更新为(2,3)，最近距离dis更新为1.5。回溯至(7,2)，如图所示，以(2,4.5)为圆心1.5为半径作圆，并不和x = 7分割超平面交割，即这里：|Q(k) - m|=|2 - 7|=5 > 1.5。至此，搜索路径回溯完。返回最近邻点(2,3)，最近距离1.5。

Kd树在维度较小时(比如20、30)，算法的查找效率很高，然而当数据维度增大(例如：K≥100)，查找效率会随着维度的增加而迅速下降。假设数据集的维数为D，一般来说要求数据的规模N满足N>>2的D次方，才能达到高效的搜索。

为了能够让Kd树满足对高维数据的索引，Jeffrey S. Beis和David G. Lowe提出了一种改进算法——Kd-tree with BBF(Best Bin First)，该算法能够实现近似K近邻的快速搜索，在保证一定查找精度的前提下使得查找速度较快。

Python 实现

Python scipy.spatial 包中封装了 KDTree 的实现

1

class KDTree(data, leafsize=10, compact_nodes=True, copy_data=False, balanced_tree=True, boxsize=None)

• Parameters:

  data array_like, shape (n,m)

  The n data points of dimension m to be indexed. This array is not copied unless this is necessary to produce a contiguous array of doubles, and so modifying this data will result in bogus results. The data are also copied if the kd-tree is built with copy_data=True.

  leafsize positive int, optional

  The number of points at which the algorithm switches over to brute-force. Default: 10.

  compact_nodes bool, optional

  If True, the kd-tree is built to shrink the hyperrectangles to the actual data range. This usually gives a more compact tree that is robust against degenerated input data and gives faster queries at the expense of longer build time. Default: True.

  **copy_data **bool, optional

  If True the data is always copied to protect the kd-tree against data corruption. Default: False.

  **balanced_tree **bool, optional

  If True, the median is used to split the hyperrectangles instead of the midpoint. This usually gives a more compact tree and faster queries at the expense of longer build time. Default: True.

  **boxsize **array_like or scalar, optional

  Apply a m-d toroidal topology to the KDTree… The topology is generated by 𝑥𝑖+𝑛𝑖𝐿𝑖 where 𝑛𝑖 are integers and 𝐿𝑖 is the boxsize along i-th dimension. The input data shall be wrapped into [0,𝐿𝑖). A ValueError is raised if any of the data is outside of this bound.

• Attributes:

  data ndarray, shape (n,m)

  The n data points of dimension m to be indexed. This array is not copied unless this is necessary to produce a contiguous array of doubles. The data are also copied if the kd-tree is built with copy_data=True.

  leafsize positive int

  The number of points at which the algorithm switches over to brute-force.

  m int

  The dimension of a single data-point.

  n int

  The number of data points.

  maxes ndarray, shape (m,)

  The maximum value in each dimension of the n data points.

  mins ndarray, shape (m,)

  The minimum value in each dimension of the n data points.

  size int

  The number of nodes in the tree.

• Methods

示例代码

1
2
3
4
5
6
7
8
9
10
11
12
13

from scipy.spatial import KDTree
if __name__ == "__main__":
    # Create a KDTree
    points = [(1, 2), (3, 4), (5, 6), (7, 8)]
    kdtree = KDTree(points)
    # Query the KDTree
    query_point = (4, 4.5)
    result = kdtree.query(query_point)
    print("The nearest neighbor of", query_point, "is", points[result[1]])

参考资料

• https://www.cnblogs.com/PythonLearner/p/12952929.html
• https://oi-wiki.org/ds/kdt/
• https://blog.csdn.net/weixin_44852067/article/details/130222731
• https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.html

文章链接：
https://www.zywvvd.com/notes/study/search/kd-tree/kd-tree/