kd樹建立的python實現

本篇博客是參考李航的《統計學習方法》第三章3.3，使用python實現kd樹的建立。kd樹本質上就是平衡二叉樹，只是注意要選擇中位數作爲節點即可

class kdNode:
    def __init__(self, data = None, depth = 0,left = None, right= None):
        self.data = data
        self.depth = depth
        self.left = left
        self.right = right
        
    def travel(self):
        if not self.data:
            return
        print self.data,self.depth
        self.left.travel()
        self.right.travel()
        
    def build_tree(self,points, depth):
        if not points:
            return
        dims = len(points[0])
        axis = depth % dims
        points.sort(key = lambda x:x[axis])
        median_index = len(points) // 2
        self.data = points[median_index]
        self.left = kdNode(None,depth+1)
        self.right = kdNode(None,depth+1)
        self.left.build_tree(points[:median_index],depth+1)
        self.right.build_tree(points[median_index+1:],depth+1)

class kdTree:
    def __init__(self, data = None,depth = 0,left = None, right = None):
        self.root = kdNode(data,depth,left,right)raise ValueError('could not find %c in %s' % (ch,str))
    #build the tree by input points
    def build_tree(self, points, depth):
        self.root.build_tree(points,depth)
    #preorder travelkd_tree.travel_tree()
    def travel_tree(self):
        self.root.travel()

通過例3.2的數據建立kd樹，使用先序遍歷查看建立結果

kd_tree = kdTree()
points = [[2,3],[5,4],[9,6],[4,7],[8,1],[7,2]]
kd_tree.build_tree(points,0)
kd_tree.travel_tree()

結果如下

[7, 2] 0
[5, 4] 1
[2, 3] 2
[4, 7] 2
[9, 6] 1
[8, 1] 2

如有錯誤，歡迎指正。