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()

高级版本

# -*- coding: utf-8 -*-
from operator import itemgetter
# kd-tree每个结点中主要包含的数据结构如下 
class KdNode(object):
    def __init__(self, median, l_child, r_child):
        self.median = median  # 位于超平面的切割点
        self.l_child = l_child  # 该结点分割超平面左子空间构成的kd-tree
        self.r_child = r_child  # 该结点分割超平面右子空间构成的kd-tree


class KdTree(object):
    def __init__(self, data):
        self.k = len(data[0])  # 假设每个点的数据维度都相同
        self.data = data    # 样本点的集合
        self.root = self.CreateNode(data, 0)  # 从第0维分量开始构建kd树,返回根节点
        self.preorder(self.root)

    def CreateNode(self, data_set, depth):
        if not data_set:  # 数据集为空
            return None

        axis = depth % self.k   #轮换选择坐标轴进行切割数据集
        # operator模块提供的itemgetter函数用于获取对象的哪些维的数据
        data_set.sort(key=itemgetter(axis))
        median = len(data_set) // 2  # 中位数分割点
        # 递归的创建kd树
        return KdNode(data_set[median],
                      self.CreateNode(data_set[:median], depth + 1),  # 创建左子树
                      self.CreateNode(data_set[median + 1:], depth + 1) ) # 创建右子树


    def visit(self, node):
        if node:
            print node.median

    # KDTree的前序遍历
    def preorder(self, root):
        if root:
            self.visit(root)
            self.preorder(root.l_child)
            self.preorder(root.r_child)


if __name__ == "__main__":
    data = [[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]]
    kd = KdTree(data)