kd-tree是KNN算法的一种实现。算法的基本思想是用多维空间中的实例点,将空间划分为多块,成二叉树形结构。划分超矩形上的实例点是树的非叶子节点,而每个超矩形内部的实例点是叶子结点。
有数据集datalist,其中的数据是Xi,每个Xi由多个特征值组成。首先将所有数据的Xi[0]找出,取得Xi[0]的中位数center,在树的根节点中保存Xi[0] == center的实例点Xi,树根的左子树递归构造Xi[0] < center的数据集,树根的右子树构造X[0] > center的数据集。同时在第二层,划分特征变为Xi[1], 划分特征随着树的深度改变,为 (d - 1) % k , k是特征的维度,d是此时划分的树深度。
import numpy as np
import matplotlib.pyplot as plot
import math
# kdtree类
class KdTree(object):
"""docstring for KdTree."""
def __str__(self):
return '{ nodes:' + str(self.nodes) + ', left:' + str(self.l) + ', right:' + str(self.r) + '}'
def __init__(self):
self.split = None
self.l = None
self.r = None
self.f = None
self.nodes = []
def createKdTree(split, datalist, k):
if datalist is None or len(datalist) == 0:
return None
split = split % k
node = KdTree()
# 求中位数
center = np.sort([a[split] for a in datalist])[int(len(datalist)/2)]
leftData = [a for a in datalist if a[split] < center]
rightData = [a for a in datalist if a[split] > center]
node.split = split
node.nodes = [a for a in datalist if a[split] == center]
node.l = createKdTree(split+1, leftData, k)
node.r = createKdTree(split+1, rightData, k)
# 设置双亲节点
if not node.l is None:
node.l.f = node
if not node.r is None:
node.r.f = node
return node
# 构建训练数据
def createData():
k = 2
datalist = None
# 构造 三类满足正态分布的类型
X_01 = np.random.randn(20) + 2
X_02 = np.random.randn(20) + 3
Y_0 = np.full(20, 1)
# 第二类
X_11 = np.random.randn(20) + 2
X_12 = 2*np.random.randn(20) + 10
Y_1 = np.full(20, 2)
# 第三类
X_21 = np.random.randn(20) + 8
X_22 = 2*np.random.randn(20) + 10
Y_2 = np.full(20, 3)
# 合并
X_1 = np.append(np.append(X_01, X_11), X_21)
X_2 = np.append(np.append(X_02, X_12), X_22)
Y = np.append(np.append(Y_0, Y_1), Y_2)
return (list(zip(X_1, X_2, Y)), k)
# 预测,返回测试集X中每个实例属于哪一类
def predict(head, X):
res = []
for x in X:
res.append(guessX(head, x))
return res
def guessX(node, x):
if x[node.split] < node.nodes[0][node.split] and not node.l is None:
return guessX(node.l, x)
elif x[node.split] > node.nodes[0][node.split] and not node.r is None:
return guessX(node.r, x)
else:
return neast(node, x)
def getDis(X1, X2):
l = len(X2)
sum = 0
for i in range(l):
sum += (X1[i] - X2[i])**2
return sum
def get2Min(nodes, x, minDis, minX):
for n in nodes:
d = getDis(n, x)
if d < minDis:
minDis = d
minX = n
return (minDis, minX)
def neast(node, x):
nodes = node.nodes
dis = []
minDis = 10000
minX = 0
minDis, minX = get2Min(nodes, x, minDis, minX)
if node.f is None:
return minX[-1]
return findY(node.f, x, minDis, minX, 0 if node.f.l == node else 1)
def findY(node, x, minDis, minX, dire):
if math.fabs(node.nodes[0][node.split] - x[node.split]) > minDis:
return minX[-1]
minDis, minX = get2Min(node.nodes, x, minDis, minX)
minDis, minX = reachAll(node.r if dire == 0 else node.l, x, minDis, minX)
if node.f is None:
# print(minX)
return minX[-1]
return findY(node.f, x, minDis, minX, 0 if node.f.l == node else 1)
def reachAll(node, x, minDis, minX):
if node is None:
return (minDis, minX)
minDis, minX = get2Min(node.nodes, x, minDis, minX)
minDis, minX = reachAll(node.l, x, minDis, minX)
minDis, minX = reachAll(node.r, x, minDis, minX)
return (minDis, minX)
if __name__ == '__main__':
# 特征 + 类别
datalist, k = createData()
head = createKdTree(0, datalist, k)
# 测试数据
X_1test = np.random.randn(5) + 2
X_2test = 2*np.random.randn(5) + 10
X = list(zip(X_1test, X_2test))
# 结果以list形式返回
res = predict(head, X)
print(res)