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
,
…
,
x
N
x_0,x_1,…,x_N
x0,x1,…,xN}
输出:n个对象的节点组成的kd树
- 初始化:以所有训练集作为样本空间,l=0;
- 构造根节点:选取样本空间内 x l x^l xl的中位数a,对应样本为 x i x_i xi,将 x i x_i xi置为根节点,左子节点的样本空间为特征 x l x^l xl小于a的子空间,右子节点的样本空间为特征 x l x^l xl大于a的子空间。
- 递归:以左右子节点作为根节点, l = ( l + 1 ) % n l=(l+1)\%n l=(l+1)%n ;重复步骤2。
- 递归终点:重复直到根节点样本空间为空,即所有样本被划分完毕。
步骤3中
l
=
(
l
+
1
)
%
n
l=(l+1)\%n
l=(l+1)%n:+1是因为要对下一个特征进行划分,取模是因为对于特征少而样本多的数据而言,划分完一轮
(
x
0
,
x
n
)
(x^0,x^n)
(x0,xn)后,可能还有样本没有划分完,由此我们可以再从特征
x
0
x^0
x0开始,直到样本均被划分。
此外,
d
e
p
t
h
%
n
=
l
depth\%n=l
depth%n=l,该结论在深度搜索kd树时有重要作用
举例:
假设有如下特征空间(3个特征,5个样本):
ka树构造算法(代码思路):
- 初始化样本空间: 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 l=0 l=0.
- 对于样本空间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 x4,由此根节点为 x 4 x_4 x4,左子节点样本空间为 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 l+=1 l+=1.
- 对于样本空间
T
1
T_1
T1,
x
(
l
=
1
)
=
{
6
,
9
}
x^{(l=1)}=\{6,9\}
x(l=1)={6,9},中位数有两个,那么我们选取其中一个,比如6(计算机特性,n/2一般都是前者,懂得都懂),对应样本为
x
0
x_0
x0,由此左子节点样本空间为
T
3
=
{
}
T_3=\{\}
T3={},右子节点样本空间为
T
4
=
{
x
1
}
T_4=\{x_1\}
T4={x1};
l
+
=
1
l+=1
l+=1;
3.1. 对于 T 3 T_3 T3={},return;对于 T 4 T_4 T4= { x 1 } \{x_1\} {x1}, x ( l = 2 ) x^{(l=2)} x(l=2)={2},中位数为2,对应样本为 x 1 x_1 x1,左右子节点样本空间均为空,return。 - 对于样本空间
T
2
T_2
T2,
x
(
l
=
1
)
=
{
1
,
3
}
x^{(l=1)}=\{1,3\}
x(l=1)={1,3},同理中位数为1,对应样本为
x
2
x_2
x2,由此左子节点样本空间为
T
5
=
{
}
T_5=\{\}
T5={},右子节点样本空间为
T
6
=
{
x
3
}
T_6=\{x_3\}
T6={x3}。
l
+
=
1
l+=1
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 x_3 x3,左右子节点样本空间均为空,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)更近?
那怎么在全局找到最近邻呢?利用贪心和回溯思想!
为了方便讨论,有如下定义和方法(认真看完!):
- 当前最近点:当前离target的距离最近的点称为当前最近点 最近距离:当前最近点与target的距离称为最近距离
- 超球体:以target为球心,最近距离为半径的球称为超球体
- 结点的超平面:在划分点上,平行划分维度的平面。比如上图,我们第一次划分的时候以第一维的4为划分点,那么 x 4 x_4 x4的超平面就是x=4的平面
- 判断超球体和超平面相交的方法:
|超平面上划分点的对应维度的值 – 球心在对应维度的值| < R(最近距离)
如上面的例子,超球体的球心:(1,6,2) ,半径:5.74, x 4 x_4 x4的超平面:x=4,划分点:(4,2,6),对应维度:1。
带入上述公式:|4-1| =3 < 5.74 因此 x 4 x_4 x4的超平面与超球体相交。
因此,我们通过dfs算法找到一个当前最近点之后,还需要讨论父节点和兄弟节点是否比当前最近点离target更近。并利用递归方法,回溯直到根节点。
(详细算法描述后文再说)
- 等等?这样的话不就需要判断所有结点了吗,和顺序搜索有什么区别?
- 错了,kd树之所以能够高效搜索,是因为如果父节点的超平面不和超球体相交的话,就不需要再讨论兄弟结点了!
- 通俗地理解就是:你有若干哥哥和若干弟弟(假设身高和年龄正相关),其中一个弟弟想和你的一个哥哥比身高,但是弟弟还没你高,也就没必要叫哥哥来比了。(转换为高维数据就是kd树的构造思想)
kd树的查找算法的完整描述:
- DFS算法:从根节点出发,递归找到一个叶子结点。
递归方法:若target当前维度值小于当前结点对应维度值,递归左子树;否则右子树。
深度和维度的关系: d e p t h % n = l depth\%n=l depth%n=l ,其中n为特征的维度(个数), depth为当前结点的深度(根的深度为1), l l l为当前深度的维度 - 设置一个STACK,用来保存DFS的路径,以便回溯。
- 先利用DFS找到一个当前最近点,无论这个点离真正的最近邻有多远。
- 对STACK的元素进行回溯:
4.1 保存好当前节点(栈顶元素)、父亲节点、兄弟节点
4.2 出栈一个
4.3 判断当前节点离target的距离是否小于最近距离:
小于:更新当前最近点和最近距离
Else :4.4
4.4 判断超球体和父节点的超平面是否相交:
相交:则从兄弟节点开始DFS并把路径入栈。
不相交:4.1 - 回溯直到栈顶元素为根节点
4.3是为了利用STACK更新最近点和最近距离,4.4是为了更新STACK,理解了这一点就不难理解上述算法了!
还有一点:4.4中需要从兄弟节点开始DFS并将递归路径压入栈,这就会产生兄弟结点之间不断地入栈的问题,因此需要设置一个标志,表示如果当前结点进行过DFS,那么兄弟节点不再进行判断。
举例:
高维不好观察,所以举二维的例子:
我们利用DFS找到当前最近点是D点,STACK=[A,B,D]
对STACK进行回溯:
- 当前节点:D,父节点:B, 兄弟节点:F,出栈一个,STACK=[A,B]
我们先判断当前节点D点是否比D更接近S(这个是程序的局限性),再判断:B点的超平面(图中红线)不和超球体相交,那么我们就不需要再判断兄弟结点F了。回溯。 - 当前节点:B,父节点:A, 兄弟节点:C,出栈一个,STACK=[A]
当前节点B不比D更近,但是A的超平面和超球体相交了,因此我们有必要讨论兄弟节点C,从C点开始DFS,路径:C->E。STACK=[A,C,E] 。回溯。 - 当前节点:E,父节点:C, 兄弟节点:G,出栈一个,STACK=[ A,C]
当前节点E比D近,更新当前最近点和最近距离。C的超平面不与超球体相交,因此不必判断G。 - 当前节点:C,父节点:A, 兄弟节点:B,出栈一个,STACK=[ A]
当前节点C不比D近,且当前节点C递归过,因此不必再判断父节点和兄弟节点。 - 当前节点: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)
运行结果:
初学者,若有错误,欢迎指正!