kdtree python代码

下面是一个简单的 Python 实现 K-D Tree 的示例代码：

import heapq

class KdNode:
    def __init__(self, point, left, right):
        self.point = point
        self.left = left
        self.right = right

class KdTree:
    def __init__(self, points):
        self.root = self.build(points)

    def build(self, points, depth=0):
        n = len(points)
        if n == 0:
            return None
        k = len(points[0])
        axis = depth % k
        sorted_points = sorted(points, key=lambda point: point[axis])
        mid = n // 2
        return KdNode(
            point=sorted_points[mid],
            left=self.build(sorted_points[:mid], depth + 1),
            right=self.build(sorted_points[mid + 1:], depth + 1)
        )

    def nearest(self, point, k=1):
        heap = []
        self.search(self.root, point, heap, k)
        return [(-dist, node.point) for dist, node in sorted(heap)]

    def search(self, node, point, heap, k):
        if node is None:
            return
        dist = sum((p - q) ** 2 for p, q in zip(node.point, point))
        if len(heap) < k:
            heapq.heappush(heap, (-dist, node))
        else:
            heapq.heappushpop(heap, (-dist, node))
        axis = len(point) % len(node.point)
        if point[axis] < node.point[axis]:
            self.search(node.left, point, heap, k)
        else:
            self.search(node.right, point, heap, k)

使用方法：

points = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]
tree = KdTree(points)
nearest = tree.nearest((6, 7), k=2)
print(nearest) # [(-2, (5, 6)), (-5, (7, 8))]

上面代码展示了如何使用 K-D Tree 寻找离一个点最近的 k 个点。在这个例子中，我们构建了一个包含 5 个点的 K-D Tree，并查找距离点 (6, 7) 最近的 2 个点。输出结果为 [(-2, (5, 6)), (-5, (7, 8))]，其中第一个元素表示距离，第二个元素表示点的坐标。