kd树的生成 @ Python

kd树的生成 @ Python

# _*_ coding:utf-8 _*_
from operator import itemgetter


class Node(object):
    def __init__(self):
        # 初始化一个节点
        self.data = []
        self.l_child = None
        self.r_child = None


class kdtree(Node):
    def create_tree(self, tree, point_list, depth=0):
        try:
            k = len(point_list[0])  # 所有的点都有相同的维度
        except IndexError as e:  
            return None
        # 轮换选择axis分割轴
        axis = depth % k

        # 对列表中的点根据axis维度排序
        point_list.sort(key=itemgetter(axis))
        median = len(point_list) // 2  # 选择中位数

        # 创建子节点和子树
        if point_list is not None:
            tree.data = point_list[median]
            tree.l_child = Node()
            tree.r_child = Node()
            self.create_tree(tree.l_child, point_list[:median], depth + 1)
            self.create_tree(tree.r_child, point_list[median + 1:], depth + 1)
            # visit a tree node

    def visit(self, tree):
        # 输入[]代表空树
        if tree is not None:
            print str(tree.data) + '\t',

    # 先序遍历
    def pre_order(self, tree):
        if tree is not None:
            self.visit(tree)
            self.pre_order(tree.l_child)
            self.pre_order(tree.r_child)

    # 中序遍历
    def in_order(self, tree):
        if tree is not None:
            self.in_order(tree.l_child)
            self.visit(tree)
            self.in_order(tree.r_child)

    # 后序遍历
    def post_order(self, tree):
        if tree is not None:
            self.post_order(tree.l_child)
            self.post_order(tree.r_child)
            self.visit(tree)


def main():
    """Example usage"""
    point_list = [(2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2)]
    t = Node()
    tree = kdtree()
    tree.create_tree(t, point_list)
    tree.pre_order(t)
    print
    tree.in_order(t)
    print 
    tree.post_order(t)


if __name__ == '__main__':
    main()