kd树与knnpython代码

以下是一个简单的KD树和KNN算法的Python代码实现：

import numpy as np

class KDTree:

    def __init__(self, data):
        self.k = data.shape[1]
        self.data = data
        self.left = None
        self.right = None
        self.split_dim = None
        self.split_value = None
        self.build()

    def build(self):
        if self.data.shape[0] == 0:
            return
        self.split_dim = np.argmax(np.var(self.data, axis=0))
        self.split_value = np.median(self.data[:, self.split_dim])
        left_data = self.data[self.data[:, self.split_dim] < self.split_value]
        right_data = self.data[self.data[:, self.split_dim] >= self.split_value]
        self.left = KDTree(left_data)
        self.right = KDTree(right_data)

    def search_knn(self, point, k):
        heap = []
        self._search_knn(point, k, heap)
        return [h[1] for h in heap]

    def _search_knn(self, point, k, heap):
        if self.data is None:
            return

        dist = np.linalg.norm(point - self.data)
        if len(heap) < k:
            heap.append((dist, self.data))
            heap.sort(key=lambda x: x[0])
        elif dist < heap[-1][0]:
            heap[-1] = (dist, self.data)
            heap.sort(key=lambda x: x[0])

        if self.left is None and self.right is None:
            return

        if point[self.split_dim] < self.split_value:
            self.left._search_knn(point, k, heap)
            if point[self.split_dim] + heap[-1][0] >= self.split_value:
                self.right._search_knn(point, k, heap)
        else:
            self.right._search_knn(point, k, heap)
            if point[self.split_dim] - heap[-1][0] < self.split_value:
                self.left._search_knn(point, k, heap)

# 使用示例
data = np.random.rand(50, 2)
tree = KDTree(data)
point = np.array([0.5, 0.5])
k = 5
knn = tree.search_knn(point, k)
print(knn)

以上代码实现了一个简单的KD树和KNN算法，可以用于查找数据集中与给定点最近的k个点。