《统计学习方法》学习之二：kd树
一、kd树模型
在使用k-means等算法时，经常需要查找最近邻节点，kd树就是一种二叉树，将特征空间进行分割，以便减小搜索时间。（具体内容可以参考李航《统计学习方法》一书）。
二、代码实现
这里实现二维平面上的kd树，可以类推到n维特征空间。
（本人代码水平有限，如有错误，还请各位大牛不吝指出）
import math
class kdTreeNode(object):
'''
kd树的节点数据结构。
'''
def __init__(self, value, x_range, y_range):
'''
:param value:kd树的节点数据，为一个长度为2的坐标信息，分别表示x坐标和y坐标。
:param x_range: 当前节点数据制约的x坐标的范围。
:param y_range: 当前节点数据制约的y坐标的范围。
self.left:左子节点。
self.right:右子节点。
self.father:父节点。
'''
self.value = value
self.x_range = x_range
self.y_range = y_range
self.left = None
self.right = None
self.father = None
class kdTree(object):
'''
kd树。
'''
def __init__(self, data, axis=0, xrange=[-100, 100], yrange=[-100, 100]):
'''
:param data:需要建树的坐标信息，是一个列表 。
:param axis: 指定开始建树时的坐标轴，0或者1。
:param xrange: x的范围。
:param yrange: y的范围。
'''
assert axis in [0, 1]
self.__data = data
self.__axis = axis
self.__xrange = xrange
self.__yrange = yrange
self.root = self.__build(self.__data, self.__axis, self.__xrange, self.__yrange)
def __build(self, data, axis, xrange, yrange):
'''
递归建立kd平衡树。
:param data: 需要建树的坐标信息，是一个列表 。
:param axis: 坐标轴，0或者1。
:param xrange: 节点制约的x坐标范围。
:param yrange: 节点制约的y坐标范围。
:return: kd树的节点，kdTreeNode
'''
assert axis in [0, 1]
if len(data) == 0:
return kdTreeNode(None, xrange, yrange)
data.sort(key=lambda s: s[axis])
pos = len(data) // 2
assert xrange[0] <= data[pos][0] <= xrange[1] and yrange[0] <= data[pos][1] <= yrange[1]
myroot = kdTreeNode(data[pos], xrange, yrange)
if axis == 0:
myroot.left = self.__build(data[0:pos:], 1-axis, [xrange[0], data[pos][0]], yrange)
myroot.right = self.__build(data[pos+1::], 1-axis, [data[pos][0], xrange[1]], yrange)
else:
myroot.left = self.__build(data[0:pos:], 1 - axis, xrange, [yrange[0], data[pos][1]])
myroot.right = self.__build(data[pos + 1::], 1 - axis, xrange, [data[pos][1], yrange[1]])
myroot.left.father = myroot
myroot.right.father = myroot
return myroot
def insert(self, data):
'''
由于主要实现的是kd树，因此为了省去不必要的麻烦，每次插入一个新节点，kd树都需要重新建立。
可以采用插入节点并调整的方法来重新恢复平衡树。
:param data: 需要插入的节点坐标。
:return: None
'''
self.__data.append(data)
self.root = self.__build(self.__data, self.__axis, self.__xrange, self.__yrange)
def search_nearest_neighbor(self, point):
'''
最近邻搜索
:param point:需要被搜索最近邻的坐标，是一个长度为2的list。
:return: 最近距离和最近邻节点的坐标。
算法采用队列保存需要搜索的节点。
'''
queue = []
node = self.root
'''搜索包含该节点的叶子节点'''
while node.value is not None:
queue.append(node)
if self.__is_containing(point, xrange=node.left.x_range, yrange=node.left.y_range):
node = node.left
else:
node = node.right
'''
由于叶子节点value为None，因此mindis（最小距离）需要计算point和该叶子节点的父节点之间的距离，
同时，最近节点为该叶子节点的父节点。
'''
mindis = math.sqrt((point[0] - queue[-1].value[0])**2 + (point[1] - queue[-1].value[1])**2)
nearest = queue[-1]
queue.append(node)
while len(queue) != 0:
node = queue.pop(0)
if node.value is not None:
'''分别对node的左右两个子节点进行搜索'''
for n in [node.left, node.right]:
'''
如果point到n的距离小于或者等于mindis，则说明n可能包含比当前最近点更近的点，因此加入队列，反之则什么都不做。
'''
dis_block = self.__calculate_distence(point, n.x_range, n.y_range)
if dis_block > mindis:
pass
else:
if dis_block not in queue:
queue.append(n)
'''
计算当前节点node和point之间的距离，如果距离小于mindis，则更新最小距离和最近节点。
'''
dis = math.sqrt((point[0] - node.value[0])**2 + (point[1] - node.value[1])**2)
if dis < mindis:
mindis = dis
nearest = node
return mindis, nearest.value
pass
def __is_containing(self, point, xrange, yrange):
'''
判断point是否包含在xrange和yrange组成的区域中。
如果point在xrange和yrange做成的区域中，包括边界，则返回True，反之则返回False。
:param point:点坐标，长度为2的list。
:param xrange: x的范围，长度为2。
:param yrange: y的范围，长度为2。
:return: bool，如果point在xrange和yrange做成的区域中，包括边界，则返回True，反之则返回False。
'''
return xrange[0] <= point[0] <= xrange[1] and yrange[0] <= point[1] <= yrange[1]
def __calculate_distence(self, point, xrange, yrange):
'''
计算point到xrange和yrange组成区域的距离。
:param point:点坐标，长度为2的list。
:param xrange: x的范围，长度为2。
:param yrange: y的范围，长度为2。
:return: float，point到xrange和yrange组成区域的距离。
'''
if xrange[0] <= point[0] <= xrange[1] and yrange[0] <= point[1] <= yrange[1]:
return 0
if point[0] < xrange[0]:
if point[1] < yrange[0]:
return math.sqrt((point[0] - xrange[0])**2 + (point[1] - yrange[0])**2)
elif yrange[0] <= point[1] <= yrange[1]:
return abs(point[0] - xrange[0])
else:
return math.sqrt((point[0] - xrange[0])**2 + (point[1] - yrange[1])**2)
elif xrange[0] <= point[0] <= xrange[1]:
if point[1] < yrange[0]:
return abs(point[1] - yrange[0])
elif yrange[0] <= point[1] <= yrange[1]:
return 0
else:
return abs(point[1] - yrange[1])
else:
if point[1] < yrange[0]:
return math.sqrt((point[0] - xrange[1])**2 + (point[1] - yrange[0])**2)
elif yrange[0] <= point[1] <= yrange[1]:
return abs(point[0] - xrange[1])
else:
return math.sqrt((point[0] - xrange[1])**2 + (point[1] - yrange[1])**2)
pass
if __name__ == '__main__':
point = []
import random
'''随机生成坐标点'''
for i in range(1000):
p = [random.uniform(0, 100), random.uniform(0, 100)]
point.append(p)
root = kdTree(point, xrange=[0, 100], yrange=[0, 100])
cnt = 0
correct = 0
x = 0.01
y = 0.01
while x < 100.0:
while y < 100.0:
mindis = 2**32-1
position = None
for i in point:
dis = math.sqrt((x - i[0])**2 + (y - i[1])**2)
if dis < mindis:
mindis = dis
position = i
m, p = root.search_nearest_neighbor([x, y])
cnt += 1
if m == mindis and position == p:
correct += 1
y += 0.01
x += 0.01
print('一共检测节点数目：', cnt)
print('正确找到最近邻节点数目：', correct)
import time
begin = time.clock()
while x < 100.0:
while y < 100.0:
mindis = 2**32-1
position = None
for i in point:
dis = math.sqrt((x - i[0])**2 + (y - i[1])**2)
if dis < mindis:
mindis = dis
position = i
y += 0.01
x += 0.01
end = time.clock()
print('线性查找最近邻所用时间：', end - begin)
begin = time.clock()
while x < 100.0:
while y < 100.0:
m, p = root.search_nearest_neighbor([x, y])
y += 0.01
x += 0.01
end = time.clock()
print('kd树查找最近邻所用时间：', end - begin)
结果如下：
一共检测节点数目： 9999
正确找到最近邻节点数目： 9999
线性查找最近邻所用时间： 7.894781509948016e-07
kd树查找最近邻所用时间： 3.9473907549739996e-07此处样本数目不大，因此在查找所用的时间上相差不大，但是依然可以看出kd树的时间小于线性查找。
三、结论
kd树是一种二叉树的应用，可以减小k-means等算法中查找最近邻所需要的时间，这对于大容量的样本数据有十分积极的意义。