kd树python实现

KD树，也称K-Dimensional Tree，是一种二叉搜索树，它能够对K维空间中的数据进行快速检索。这里提供一个Python实现的示例代码：

import numpy as np

class KDTree:
    def __init__(self, X):
        self.k = X.shape[1]
        self.tree = self.build_tree(X)

    class Node:
        def __init__(self, data, dim, left, right):
            self.data = data
            self.dim = dim
            self.left = left
            self.right = right

    def build_tree(self, X):
        if X.shape[0] == 0:
            return None

        dim = np.argmax(np.var(X, axis=0))
        sorted_idx = np.argsort(X[:, dim])
        mid = X.shape[0] // 2

        left = self.build_tree(X[sorted_idx[:mid]])
        right = self.build_tree(X[sorted_idx[mid+1:]])

        return KDTree.Node(X[sorted_idx[mid]], dim, left, right)

    def search(self, x):
        def helper(node, x, best_dist, best_node):
            if node is None:
                return best_node, best_dist

            dist = np.sum((node.data - x) ** 2)
            if dist < best_dist:
                best_dist = dist
                best_node = node

            if x[node.dim] < node.data[node.dim]:
                best_node, best_dist = helper(node.left, x, best_dist, best_node)
                if x[node.dim] + np.sqrt(best_dist) > node.data[node.dim]:
                    best_node, best_dist = helper(node.right, x, best_dist, best_node)
            else:
                best_node, best_dist = helper(node.right, x, best_dist, best_node)
                if x[node.dim] - np.sqrt(best_dist) < node.data[node.dim]:
                    best_node, best_dist = helper(node.left, x, best_dist, best_node)

            return best_node, best_dist

        return helper(self.tree, x, np.inf, None)

代码中的`KDTree`类实现了KD树的构建和搜索功能。在初始化时，传入数据`X`，并根据方差最大的维度进行划分，递归构建KD树。搜索时，从根节点开始递归地遍历左右子树，更新最近邻节点和距离。具体实现过程详见代码注释。