import math
import numpy as np
class KD_node:
def __init__(self, point=None, split=None, LL = None, RR = None):
self.point = point
self.split = split
self.left = LL
self.right = RR
def createKDTree(root, data_list):
LEN = len(data_list)
if LEN == 0:
return
dimension = len(data_list[0])
max_var = 0
split = 0;
for i in range(dimension):
ll = []
for t in data_list:
ll.append(t[i])
var = computeVariance(ll)
if var > max_var:
max_var = var
split = i
data_list.sort(key=lambda x: x[split])
point = data_list[int(LEN / 2)]
root = KD_node(point, split)
root.left = createKDTree(root.left, data_list[0:int(LEN / 2)])
root.right = createKDTree(root.right, data_list[int(LEN / 2 + 1):LEN])
return root
def computeVariance(arrayList):
LEN = len(arrayList)
array = np.array(arrayList)
sum1 = array.sum()
array2 = array * array
sum2 = array2.sum()
mean = sum1 / LEN
variance = sum2 / LEN - mean**2
return variance
def findNN(root, query):
NN = root.point
min_dist = computeDist(query, NN)
nodeList = []
temp_root = root
while temp_root:
nodeList.append(temp_root)
dd = computeDist(query, temp_root.point)
if min_dist > dd:
NN = temp_root.point
min_dist = dd
ss = temp_root.split
if query[ss] <= temp_root.point[ss]:
temp_root = temp_root.left
else:
temp_root = temp_root.right
while nodeList:
back_point = nodeList.pop()
ss = back_point.split
print( "back.point = ", back_point.point )
if abs(query[ss] - back_point.point[ss]) < min_dist:
if query[ss] <= back_point.point[ss]:
temp_root = back_point.right
else:
temp_root = back_point.left
if temp_root:
nodeList.append(temp_root)
curDist = computeDist(query, temp_root.point)
if min_dist > curDist:
min_dist = curDist
NN = temp_root.point
return NN, min_dist
def computeDist(pt1, pt2):
sum = 0.0
for i in range(len(pt1)):
sum = sum + (pt1[i] - pt2[i]) * (pt1[i] - pt2[i])
return math.sqrt(sum)
if __name__=="__main__":
root=KD_node()
l=[[2,3],[5,4],[9,6],[4,7],[8,1],[7,2]]
t1=createKDTree(root,l)
print(findNN(t1,[2.1,3.1]))