《统计学习方法》KD树的构建与查找Python实现
关于KD树的构建和查找的理论可参考《统计学习方法》第三章以及这篇博文https://zhuanlan.zhihu.com/p/23966698，
import os
import csv
import numpy as np
import string
import pandas as pd
import operator
import re as re
import time
import datetime
class KD_node():
def __init__(self, vector, label, depth, dimension, l_node = None, r_node = None):
self.val = vector
self.label = label
self.depth = depth
self.dimension = dimension
self.l_node = l_node
self.r_node = r_node
def add_lnode(self, new_node):
self.l_node = new_node
def add_rnode(self, new_node):
self.r_node = new_node
def print_preorder(self):
print self.val
if self.l_node:
self.l_node.print_preorder()
if self.r_node:
self.r_node.print_preorder()
def search(self, fea_vec, L, k):
if self.val[self.dimension] < fea_vec[self.dimension]:#目标点在当前结点左边
if self.l_node != None:
L = self.l_node.search(fea_vec, L, k) #若还有左结点则搜索左结点
L = self.insertL(fea_vec, L, k) #对当前结点执行插入操作（包含了是否插入条件检测）
if len(L) < k: #若L仍然未满，则不用考虑，直接搜索右子树
L = self.r_node.search(fea_vec, L, k)
else if (self.val[self.dimension] - fea_vec[self.dimension])**2 < L[-1][1]: #判断右子树是否可能存在能插入L的结点，判断方法是求目标点到切割超平面的距离，此距离也就是切割维度上，切割点与目标点的距离，并将此距离与L内最大距离比较，若小于L内最大距离，则说明在切割超平面的另一边仍有可能存在能插入L的结点，我认为这一步，这个判断条件也是KD树提高搜索效率的关键所在
L = self.r_node.search(fea_vec, L, k)
else: #目标点在当前结点右边，跟上边的情况是对称的，就不写那么多注释了
if self.r_node != None:
L = self.r_node.search(fea_vec, L, k)
L = self.insertL(fea_vec, L, k)
if len(L) < k:
L = self.l_node.search(fea_vec, L, k)
else if (self.val[self.dimension] - fea_vec[self.dimension])**2 < L[-1][1]:
L = self.l_node.search(fea_vec, L, k)
return L
def insertL(self, fea_vec, L, k): #插入操作，检查当前结点是否可以插入L，若可以则插入，否则不做操作
distance = sum((fea_vec - self.val)**2)
if len(L) < k: #如果L还没满，则直接插入
L.append([self, distance])
L = L[L[:, 1].argsort()] #插入后排序，对Python不太熟，急于实现算法，用的效率比较低的方式，见谅。。。
else if distance < L[-1][1]： #若L已满，判断L中最大距离是否比当前结点距离更大
L[-1] = [self, distance]
L = L[L[:, 1].argsort()]
return L
def construct(data, depth):
if len(data) == 0:
return None
dimension_sum = np.shape(data)[1] - 1 #特征维度,减1是因为label也占了一列
dimension = depth % dimension_sum + 1 #切割轴维度
data = data[data[:, dimension].argsort()]
node_data = data[len(data)/2]
new_node = KD_node(node_data[1:], node_data[0], depth, dimension)
new_node.l_node = construct(data[:len(data)/2], depth + 1)
new_node.r_node = construct(data[len(data)/2 + 1:], depth + 1)
return new_node
class KD_Tree(): #与一般的二叉排序树结构差不多，区别在于KD树需要反复使用各个维度来比较以构造二叉树
# def __init__(self, data_root):
# data_file = pd.read_csv(data_root, header = None)
# self.data = np.array(data_file)[:, 2:]
# self.label = np.array(data_file)[:, 1]
def __init__(self, data):
self.data = data[:,:2]
self.label = data[:, 2]
def constructor(self):
data = np.c_[self.label, self.data] #合并特征向量和label
data = data[data[:, 1].argsort()] #按第一维排序（第0维是label）
self.root = construct(data, 0)
def print_preorder(self):
print 'preorder'
self.root.print_preorder()
def predict(self, fea_vec, k = 3):
L = [] #存储当前检测到的距离最近的K个点，初始为空，L中的元素为[KD_node, distance]
self.root(fea_vec, L, k)
label_dict = {}
for each_sample in L:
label_dict[each_sample.label] += 1
print(sorted(label_dict,key=lambda x:label_dict[x])[-1]) #打印出k个最近邻结点中个数最多的类别