树方法

目录

• 树方法
  • kd-tree

kd-tree

kd-tree (k dimensional tree )是树方法的经典算法，其是二分搜索树在多维空间的推广。二分搜索树检索迅速的原因是规定将数据中大于当前节点数据的方在一侧(比如右子树)，而不小于的放在另一侧(比如左子树),这样检索数据时，即可获得logn的速度。kd-tree 也类似,也是二叉搜索树，只不过其中的一维数据变成了n维数据。如同x-y轴坐标系将二维空间分成四个部分一样， n维空间也可这样划分。然后，规定将大于分割点的数据放在某一侧(比如右子树)，不小于分割点的数据放在另一侧(比如左子树)。 kd-tree 适用\(N>>2^k\) 的情形，其中N为数据数量。不过实际生产中，k不能太大，一般10维左右时，效果不错，最好不要超过20维。

from collections import namedtuple
from operator import itemgetter
from pprint import pformat

class Node(namedtuple("Node", "location left_child right_child")):
    def __repr__(self):
        return pformat(tuple(self))

def kdtree(point_list, depth: int = 0):
    if not point_list:
        return None

    k = len(point_list[0])  # assumes all points have the same dimension
    # Select axis based on depth so that axis cycles through all valid values
    axis = depth % k

    # Sort point list by axis and choose median as pivot element
    point_list.sort(key=itemgetter(axis))
    median = len(point_list) // 2

    # Create node and construct subtrees
    return Node(
        location=point_list[median],
        left_child=kdtree(point_list[:median], depth + 1),
        right_child=kdtree(point_list[median + 1 :], depth + 1),
    )

def main():
    """Example usage"""
    point_list = [(7, 2), (5, 4), (9, 6), (4, 7), (8, 1), (2, 3)]
    tree = kdtree(point_list)
    print(tree)

if __name__ == "__main__":
    main()