KD树搜索

KD树的创建和KD树的搜索实现：

import numpy as np
def createTree(dataset,depth):
    """
    dataset为数据集，depth为深度
    """
    treeNode={}
    m=np.shape(dataset)[0]
    if m==0:
        return None
    else:
        m,n=np.shape(dataset)    
        split=(depth%n)
        dataset=sorted(dataset,key=lambda x:x[split])
        median=m//2       #获取中位数
        treeNode['split']=split   #设置按照哪一维度划分
        treeNode['median']=dataset[median]  #设置中位数组
        depth+=1
        treeNode['left']=createTree(dataset[0:median],depth)  #创建左子树
        treeNode['right']=createTree(dataset[median+1:],depth) #创建右子树
        return treeNode

        
def searchTree(treeNode,data):
    """
    data为单独点测试集，为了搜索到最近邻
    """
    #n=np.shape(data)[0]
    if treeNode is None:
        return [0]*len(data),float('inf')
    split=treeNode['split']
    median_point=treeNode['median']
    if data[split]<=median_point[split]:   #递归到叶子节点
        nearestPoint,nearestDistance=searchTree(treeNode['left'],data)
    else:
        nearestPoint,nearestDistance=searchTree(treeNode['right'],data)
    nowDistance=np.linalg.norm(data-median_point)  #计算第二范数的欧式距离
    if nowDistance<nearestDistance:
        nearestDistance=nowDistance
        nearestPoint=median_point.copy()
    splitDistance=abs(data[split]-median_point[split]) #计算维度之间的距离，判断是否与目标点与当前最近点所够成的圆内
    if splitDistance>nearestDistance:  #维度距离大于最小点
        return nearestPoint,nearestDistance
    else:   #回溯，寻找与圆相交的点，同时向结点分类相反方向搜索结点，计算两点的距离，获取距离最小值。
        if data[split]<=median_point[split]:
            nextTree=treeNode['right']
            #nearPoint,nearDistance=searchTree(treeNode['right'],data)
        else:
            nextTree=treeNode['left']
        nearPoint,nearDistance=searchTree(nextTree,data)
        if nearDistance<nearestDistance:
            nearestDistance=nearDistance
            nearestPoint=nearPoint.copy()
        return nearestPoint,nearestDistance


dataset=np.array([[1,2],[3,4],[5,4],[7,2],[6,3],[8,7]])
tree=createTree(dataset,0)
tree
#dataset[2]
m,n=searchTree(tree,[1,5])
m,n

Out[85]:
(array([3, 4]), 2.23606797749979)