python kd树 搜索

Posted 2020-10-24 Alice_鹿_Bambi

用python实现kd树的构造和搜索目标点最近邻的过程

　　kd树就是一种对k维空间中的实例点进行存储以便对其进行快速检索的树形数据结构，可以运用在k近邻法中，实现快速k近邻搜索。构造kd树相当于不断地用垂直于坐标轴的超平面将k维空间切分，依次选择坐标轴对空间进行切分，选择训练实例点在选定坐标轴上的中位数为切分点。具体kd树的原理可以参考kd树的原理。

　　代码是参考《统计学习方法》k近邻 kd树的python实现得到

　　首先创建一个类，用于表示树的节点，包括：该节点的值，用于划分左右子树的切分轴，左子树，右子树

class decisionnode:
    def __init__(self,value=None,col=None,rb=None,lb=None):
        self.value=value
        self.col=col
        self.rb=rb
        self.lb=lb

　　切分点为坐标轴上的中值，下面代码求得一个序列的中值

def median(x):
    n=len(x)
    x=list(x)
    x_order=sorted(x)
    return x_order[n//2],x.index(x_order[n//2])

然后就可以构造一颗kd树，左子树小于切分点，右子树大于切分点，data是输入的数据

def buildtree(x,j=0):
    rb=[]
    lb=[]
    m,n=x.shape
    if m==0: return None
    edge,row=median(x[:,j].copy())
    for i in range(m):
        if x[i][j]>edge: 
            rb.append(i)
        if x[i][j]<edge:
            lb.append(i)
    rb_x=x[rb,:]
    lb_x=x[lb,:]
    rightBranch=buildtree(rb_x,(j+1)%n)
    leftBranch=buildtree(lb_x,(j+1)%n)
    return decisionnode(x[row,:],j,rightBranch,leftBranch)

　　接下来是树的搜索过程，可以用下图表示树的搜索过程，具体过程可以参考kd树的原理。

　　

　　代码如下：

#搜索树：nearestPoint,nearestValue均为全局变量
def traveltree(node,point):
    global nearestPoint,nearestValue
    if node==None: return 
    print(node.value)
    print(\'---\')
    col=node.col
    if point[col]>node.value[col]:
        traveltree(node.rb,point)
    if point[col]<node.value[col]:
        traveltree(node.lb,point)
    dis=dist(node.value,point)
    print(dis)
    if dis<nearestValue:
        nearestPoint=node
        nearestValue=dis
        #print(\'nearestPoint,nearestValue\' % (nearestPoint,nearestValue))
    if node.rb!=None or node.lb!=None:
        if abs(point[node.col] - node.value[node.col]) < nearestValue:
            if point[node.col]<node.value[node.col]:
                traveltree(node.rb,point)
            if point[node.col]>node.value[node.col]:
                traveltree(node.lb,point)
        
def searchtree(tree,aim):
    global nearestPoint,nearestValue
    #nearestPoint=None
    nearestValue=float(\'inf\')
    traveltree(tree,aim)
    return nearestPoint
        
    
def dist(x1, x2): #欧式距离的计算  
    return ((np.array(x1) - np.array(x2)) ** 2).sum() ** 0.5

 完整代码在此处取

 1 import numpy as np
 2 from numpy import array
 3 class decisionnode:
 4     def __init__(self,value=None,col=None,rb=None,lb=None):
 5         self.value=value
 6         self.col=col
 7         self.rb=rb
 8         self.lb=lb
 9         
10 #读取数据并将数据转换为矩阵形式        
11 def readdata(filename):    
12     data=open(filename).readlines()
13     x=[]
14     for line in data:
15         line=line.strip().split(\'\\t\')
16         x_i=[]
17         for num in line:
18             num=float(num)
19             x_i.append(num)
20         x.append(x_i)
21     x=array(x)
22     return x
23 
24 #求序列的中值    
25 def median(x):
26     n=len(x)
27     x=list(x)
28     x_order=sorted(x)
29     return x_order[n//2],x.index(x_order[n//2])
30 
31 #以j列的中值划分数据，左小右大，j=节点深度%列数    
32 def buildtree(x,j=0):
33     rb=[]
34     lb=[]
35     m,n=x.shape
36     if m==0: return None
37     edge,row=median(x[:,j].copy())
38     for i in range(m):
39         if x[i][j]>edge: 
40             rb.append(i)
41         if x[i][j]<edge:
42             lb.append(i)
43     rb_x=x[rb,:]
44     lb_x=x[lb,:]
45     rightBranch=buildtree(rb_x,(j+1)%n)
46     leftBranch=buildtree(lb_x,(j+1)%n)
47     return decisionnode(x[row,:],j,rightBranch,leftBranch)
48 
49 #搜索树：nearestPoint,nearestValue均为全局变量
50 def traveltree(node,point):
51     global nearestPoint,nearestValue
52     if node==None: return 
53     print(node.value)
54     print(\'---\')
55     col=node.col
56     if point[col]>node.value[col]:
57         traveltree(node.rb,point)
58     if point[col]<node.value[col]:
59         traveltree(node.lb,point)
60     dis=dist(node.value,point)
61     print(dis)
62     if dis<nearestValue:
63         nearestPoint=node
64         nearestValue=dis
65         #print(\'nearestPoint,nearestValue\' % (nearestPoint,nearestValue))
66     if node.rb!=None or node.lb!=None:
67         if abs(point[node.col] - node.value[node.col]) < nearestValue:
68             if point[node.col]<node.value[node.col]:
69                 traveltree(node.rb,point)
70             if point[node.col]>node.value[node.col]:
71                 traveltree(node.lb,point)
72         
73 def searchtree(tree,aim):
74     global nearestPoint,nearestValue
75     #nearestPoint=None
76     nearestValue=float(\'inf\')
77     traveltree(tree,aim)
78     return nearestPoint
79         
80     
81 def dist(x1, x2): #欧式距离的计算  
82     return ((np.array(x1) - np.array(x2)) ** 2).sum() ** 0.5  

kdtree