KNN算法之KD树（K-dimension Tree）实现 K近邻查询

KD树是一种分割k维数据空间的数据结构，主要应用于多维空间关键数据的搜索，如范围搜索和最近邻搜索。

KD树使用了分治的思想，对比二叉搜索树（BST），KD树解决的是多维空间内的最近点（K近点）问题。（思想与之前见过的最近点对问题很相似，将所有点分为两边，对于可能横跨划分线的点对再进一步讨论）

KD树用来优化KNN算法中的查询复杂度。

一、建树

建立KDtree，主要有两步操作：选择合适的分割维度，选择中值节点作为分割节点。

分割维度的选择遵循的原则是，选择范围最大的纬度，也即是方差最大的纬度作为分割维度，为什么方差最大的适合作为特征呢？

因为方差大，数据相对“分散”，选取该特征来对数据集进行分割，数据散得更“开”一些。

分割节点的选择原则是，将这一维度的数据进行排序，选择正中间的节点作为分割节点，确保节点左边的点的维度值小于节点的维度值，节点右边的点的维度值大于节点的维度值。

这两步步骤影响搜索效率，非常关键。

二、搜索K近点

需要的数据结构：最大堆（此处我对距离取负从而用最小堆实现的最大堆，因为python的heapq模块只有最小堆）、堆栈（用列表实现）

a.利用二叉搜索找到叶子节点并将搜索的结点路径压入堆栈stack中

b.通过保存在堆栈中的搜索路径回溯，直至堆栈中没有结点了

对于b步骤，需要区分叶子结点和非叶结点：

1、叶子结点：

叶子结点:计算与目标点的距离。若候选堆中不足K个点，将叶子结点加入候选堆中；如果K个点够了，判断是否比候选堆中距离最小的结点（因为距离取了相反数）还要大， 说明应当加入候选堆；

2、非叶结点：

对于非叶结点，处理步骤和叶子结点差不多，只是需要额外考虑以目标点为圆心，最大堆的堆顶元素为半径的超球体是否和划分当前空间的超平面相交，如果相交说明未访问的另一边的空间有可能包含比当前已有的K个近点更近的点，需要搜索另一边的空间；此外，当候选堆中没有K个点，那么不管有没有相交，都应当搜索未访问的另一边空间，因为候选堆的点不够K个。

步骤：计算与目标点的距离
1、若不足K个点，将结点加入候选堆中；
如果K个点够了，判断是否比候选堆中距离最小的结点（因为距离取了相反数）还要大。

2、判断候选堆中的最小距离是否小于Xi离当前超平面的距离（即是否需要判断未访问的另一边要不要搜索）当然如果不足K个点，虽然超平面不相交，依旧要搜索另一边，直到找到叶子结点，并且把路径加入回溯栈中。

三、预测

KNN通常用来分类或者回归问题，笔者已经封装好了两种预测的方法。

python代码实现：

  1 import heapq
  2 class KDNode(object):
  3     def __init__(self,feature=None,father=None,left=None,right=None,split=None):
  4         self.feature=feature # dimension index (按第几个维度的特征进行划分的)
  5         self.father=father  #并没有用到
  6         self.left=left
  7         self.right=right
  8         self.split=split # X value and Y value （元组，包含特征X和真实值Y）
  9
 10 class KDTree(object):
 11     def __init__(self):
 12         self.root=KDNode()
 13         pass
 14
 15     def _get_variance(self,X,row_indexes,feature_index):
 16         # X (2D list): samples * dimension
 17         # row_indexes (1D list): choose which row can be calculated
 18         # feature_index (int): calculate which dimension
 19         n = len(row_indexes)
 20         sum1 = 0
 21         sum2 = 0
 22         for id in row_indexes:
 23             sum1 = sum1 + X[id][feature_index]
 24             sum2 = sum2 + X[id][feature_index]**2
 25
 26         return sum2/n - (sum1/n)**2
 27
 28     def _get_max_variance_feature(self,X,row_indexes):
 29         mx_var = -1
 30         dim_index = -1
 31         for dim in range(len(X[0])):
 32             dim_var = self._get_variance(X,row_indexes,dim)
 33             if dim_var>mx_var:
 34                 mx_var=dim_var
 35                 dim_index=dim
 36         # return max variance feature index (int)
 37         return dim_index
 38
 39     def _get_median_index(self,X,row_indexes,feature_index):
 40         median_index =  len(row_indexes)//2
 41         select_X = [(idx,X[idx][feature_index]) for idx in row_indexes]
 42         sorted_X = select_X
 43         sorted(sorted_X,key= lambda x:x[1])
 44         #return median index in feature_index dimension (int)
 45         return sorted_X[median_index][0]
 46
 47     def _split_feature(self,X,row_indexes,feature_index,median_index):
 48         left_ids = []
 49         right_ids = []
 50         median_val = X[median_index][feature_index]
 51         for id in row_indexes:
 52             if id==median_index:
 53                 continue
 54             val = X[id][feature_index]
 55             if val < median_val:
 56                 left_ids.append(id)
 57             else:
 58                 right_ids.append(id)
 59         # return (left points index and right points index)(list,list)
 60         # 把当前的样本按feature维度进行划分为两份
 61         return left_ids, right_ids
 62
 63     def build_tree(self,X,Y):
 64         row_indexes =[i for i in range(len(X))]
 65         node =self.root
 66         queue = [(node,row_indexes)]
 67         # BFS创建KD树
 68         while queue:
 69             root,ids = queue.pop(0)
 70             if len(ids)==1:
 71                 root.feature = 0 #如果是叶子结点，维度赋0
 72                 root.split = (X[ids[0]],Y[ids[0]])
 73                 continue
 74             # 选取方差最大的特征维度划分,取样本的中位数作为median
 75             feature_index = self._get_max_variance_feature(X,ids)
 76             median_index = self._get_median_index(X,ids,feature_index)
 77             left_ids,right_ids = self._split_feature(X,ids,feature_index,median_index)
 78             root.feature = feature_index
 79             root.split = (X[median_index],Y[median_index])
 80             if left_ids:
 81                 root.left = KDNode()
 82                 root.left.father = root
 83                 queue.append((root.left,left_ids))
 84             if right_ids:
 85                 root.right = KDNode()
 86                 root.right.father = root
 87                 queue.append((root.right,right_ids))
 88
 89     def _get_distance(self,Xi,node,p=2):
 90         # p=2  default Euclidean distance
 91         nx = node.split[0]
 92         dist = 0
 93         for i in range(len(Xi)):
 94             dist=dist + (abs(Xi[i]-nx[i])**p)
 95         dist = dist**(1/p)
 96         return dist
 97
 98     def _get_hyperplane_distance(self,Xi,node):
 99         xx = node.split[0]
100         dist = abs(Xi[node.feature] - xx[node.feature])
101         return dist
102
103     def _is_leaf(self,node):
104         if node.left is None and node.right is None:
105             return True
106         else:
107             return False
108
109     def get_nearest_neighbour(self,Xi,K=1):
110         search_paths = []
111         max_heap = [] #use min heap achieve max heap （因为python只有最小堆）
112         priority_num = 0  # remove same distance
113         heapq.heappush(max_heap,(float('-inf'),priority_num,None))
114         priority_num +=1
115         node = self.root
116         # 找到离Xi最近的叶子结点
117         while node is not None:
118             search_paths.append(node)
119             if Xi[node.feature] < node.split[0][node.feature]:
120                 node = node.left
121             else:
122                 node = node.right
123
124         while search_paths:
125             now = search_paths.pop()
126             #  叶子结点:计算与Xi的距离，若不足K个点，将叶子结点加入候选堆中；
127             #  如果K个点够了，判断是否比候选堆中距离最小的结点（因为距离取了相反数）还要大，
128             #  说明应当加入候选堆；
129             if self._is_leaf(now):
130                 dist = self._get_distance(Xi,now)
131                 dist = -dist
132                 mini_dist = max_heap[0][0]
133                 if len(max_heap) < K :
134                     heapq.heappush(max_heap,(dist,priority_num,now))
135                     priority_num+=1
136                 elif dist > mini_dist:
137                     _ = heapq.heappop(max_heap)
138                     heapq.heappush(max_heap,(dist,priority_num,now))
139                     priority_num+=1
140             #  非叶结点:计算与Xi的距离
141             #  1、若不足K个点，将结点加入候选堆中；
142             #   如果K个点够了，判断是否比候选堆中距离最小的结点（因为距离取了相反数）还要大，
143             #  2、判断候选堆中的最小距离是否小于Xi离当前超平面的距离（即是否需要判断另一边要不要搜索）
144             #    当然如果不足K个点，虽然超平面不相交，依旧要搜索另一边，
145             #    直到找到叶子结点，并且把路径加入回溯栈中
146             else :
147                 dist = self._get_distance(Xi, now)
148                 dist = -dist
149                 mini_dist = max_heap[0][0]
150                 if len(max_heap) < K:
151                     heapq.heappush(max_heap, (dist, priority_num, now))
152                     priority_num += 1
153                 elif dist > mini_dist:
154                     _ = heapq.heappop(max_heap)
155                     heapq.heappush(max_heap, (dist, priority_num, now))
156                     priority_num += 1
157
158                 mini_dist = max_heap[0][0]
159                 if len(max_heap)<K or -(self._get_hyperplane_distance(Xi,now)) > mini_dist:
160                     # search another child tree
161                     if Xi[now.feature] >= now.split[0][now.feature]:
162                         child_node = now.left
163                     else:
164                         child_node = now.right
165                     # record path until find child leaf node
166                     while child_node is not None:
167                         search_paths.append(child_node)
168                         if Xi[child_node.feature] < child_node.split[0][child_node.feature]:
169                             child_node = child_node.left
170                         else:
171                             child_node = child_node.right
172         return max_heap
173
174     def predict_classification(self,Xi,K=1):
175         # 多分类问题预测
176         y =self.get_nearest_neighbour(Xi,K)
177         mp = {}
178         for i in y:
179             if i[2].split[1] in mp:
180                 mp[i[2].split[1]]+=1
181             else:
182                 mp[i[2].split[1]]=1
183         pre_y = -1
184         max_cnt =-1
185         for k,v in mp.items():
186             if v>max_cnt:
187                 max_cnt=v
188                 pre_y=k
189         return pre_y
190
191     def predict_regression(self,Xi,K=1):
192         #回归问题预测
193         pre_y = self.get_nearest_neighbour(Xi,K)
194         return sum([i[2].split[1] for i in pre_y])/K
195
196
197 # t =KDTree()
198 # xx = [[3,3],[1,2],[5,6],[999,999],[5,5]]
199 # z = [1,0,1,1,1]
200 # t.build_tree(xx,z)
201 # y=t.predict_regression([4,5.5],K=5)
202 # print(y)

参考资料：

《统计学习方法》——李航著

https://blog.csdn.net/qq_32478489/article/details/82972391?utm_medium=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromMachineLearnPai2-1.channel_param&depth_1-utm_source=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromMachineLearnPai2-1.channel_param

https://zhuanlan.zhihu.com/p/45346117

https://www.cnblogs.com/xingzhensun/p/9693362.html