KDTree

 

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]))