不使用sklearn完成对KD树的划分

import numpy as np # 节点对象 class KDTreeNode: def __init__(self, data,index, dim, right=None, left=None): self.data = data self.index = index self.dim = dim self.right = right self.left = left def __str__(self): return f'特征值:{self.data}\n 下标 ：{self.index}\n 划分的维度: {self.dim}' # 欧式距离的计算 def dist(x1, x2): return ((np.array(x1) - np.array(x2)) ** 2).sum() ** 0.5 class MyKDTree: def __init__(self): self.nearestValue = None self.nearestPoint = None self.tree=None self.data=None self.target=None def kdtree_build(self,X, y, dim): if X.size == 0: return None nidx = np.argsort(X, axis=0)[:, dim] # 取中位数 node_idx = X.shape[0] // 2 # 中位数以左为左子树 left_index = nidx[: node_idx] # 1, 0, 2 # 中位数以右为右子树 right_index = nidx[node_idx + 1:] # 5, 4 # 中位数为当前节点 cur_index = nidx[node_idx] # 3 # 创建当前节点 cur_node = KDTreeNode(X[cur_index], cur_index, dim) # print(f'当前节点:\n{cur_node}, index: {cur_index}') # print(f'左子树节点:\n {X[left_index]}') # print(f'右子树节点:\n {X[right_index]}') # 划分的维度要往后挪1位 dim = (dim + 1) % X.shape[1] # 左子树 cur_node.left=self.kdtree_build(X[left_index], y[left_index], dim) # 右子树 cur_node.right=self.kdtree_build(X[right_index], y[right_index], dim) return cur_node def fit(self,X,y): self.data=X self.target=y self.tree=self.kdtree_build(X,y,0) def predict(self,x): index=[] for i in x: index.append(self.search(self.tree,i)) return self.target[index] def search(self, tree, x): self.nearestPoint = None # 保存最近的点 self.nearestValue = 0 # 保存最近的值 def travel(node, depth=0): # 递归搜索 if node is not None: # 递归终止条件 n = len(x) # 特征数 axis = depth % n # 计算轴 if x[axis] < node.data[axis]: # 如果数据小于结点，则往左结点找 travel(node.left, depth + 1) else: travel(node.right, depth + 1) # 以下是递归完毕后，往父结点方向回朔 distNodeAndX = dist(x, node.data) # 目标和节点的距离判断 if self.nearestPoint is None: # 确定当前点，更新最近的点和最近的值 self.nearestPoint = node.index self.nearestValue = distNodeAndX elif self.nearestValue > distNodeAndX: self.nearestPoint = node.index self.nearestValue = distNodeAndX # print(node.data, depth, self.nearestValue, node.data[axis], x[axis]) if abs(x[axis] - node.data[axis]) <= self.nearestValue: # 确定是否需要去子节点的区域去找（圆的判断） if x[axis] < node.data[axis]: travel(node.left, depth + 1) else: travel(node.right, depth + 1) travel(tree) return self.nearestPoint