kdtree.py

# -*- coding: utf-8 -*-
"""Kdtree.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1iJ_mzsq0z_3coWmWdD8yOfrPED5uYT0n

#Library
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import nltk
import re
from numpy import random
import pprint
import queue as que
from queue import LifoQueue
import operator
import collections
import time

#reference kd-tree split value to github: https://github.com/HolyChen/KdTree/blob/master/kdtree.py

"""#Helper Function
* Euclide_Distance
* data_pre_balance
"""

def euclide_distance(a,b):
    point1 = a
    point2 = b

    dist = 0.0

    dist = pow(point1.x - point2.x,2) + pow(point1.y - point2.y,2) + pow(point1.z - point2.z,2)

    dist = np.sqrt(dist)
    
    return dist

def euclide_distance_split(a,b):

    dist = 0.0
    a = np.squeeze(a)
    b = np.squeeze(b)
    for i in range(len(a)):
        dist += pow(a[i] - b[i],2)

    dist = np.sqrt(dist)
    
    return dist

#data_pre_balance for classic kdtree
def preprocess(points,level,begin,end):
    if begin >= end:
       return None
    check = level % 3
    if check == 0:
       points[begin : end] = sorted(points[begin : end], key=operator.attrgetter('x'))
    elif check == 1:
       points[begin : end] = sorted(points[begin : end], key=operator.attrgetter('y'))
    else:
       points[begin : end] = sorted(points[begin : end], key=operator.attrgetter('z'))
    mid = int((begin + end)/2)
    preprocess(points,level + 1,begin,mid)
    preprocess(points,level + 1,mid + 1,end)

"""#Kd Tree Definition (Classic)
* class Point definition
* class Node definition
* Helper FunctionRec for class Kd Tree
* class KdTree definition

class Point definition
"""

class Point:
      def __init__(self,*args):
          if len(args) == 0:
             self.x = 0.0
             self.y = 0.0
             self.z = 0.0
          elif len(args) == 1 and isinstance(args[0],Point):
             self.x = args[0].x
             self.y = args[0].y
             self.z = args[0].z
          else:
             self.x = args[0]
             self.y = args[1]
             self.z = args[2]

      def set_x(self,x):
          self.x = x
      
      def set_y(self,y):
          self.y = y
      
      def set_z(self,z):
          self.z = z

"""class Node definition"""

class Node_normal:
      def __init__(self,point):
          self.data = point
          self.left = None
          self.right = None

"""Helper FunctionRec for class KdTree"""

#==============================================================================
def checkPoint(a,b):
    if a.x != b.x:
       return False
    if a.y != b.y:
       return False
    if a.z != b.z:
       return False
    return True
#==============================================================================

#==============================================================================
def insertRec(root,data,level):
    if root is None:
       return Node_normal(data)
    check = level % 3
    root_cor = 0.0
    data_cor = 0.0

    #gen_data of root_cor and data_cor    
    if check == 0:
       root_cor = root.data.x
       data_cor = data.x
    elif check == 1:
       root_cor = root.data.y
       data_cor = data.y
    else:
       root_cor = root.data.z
       data_cor = data.z

    if root_cor > data_cor:
       root.left = insertRec(root.left,data,level + 1)
    else:
       root.right = insertRec(root.right,data,level + 1)
    return root
#==============================================================================

#==============================================================================
def searchRec(root,point,level):
    if root is None:
       return None
    check = level % 3
    root_cor = 0.0
    data_cor = 0.0

    if checkPoint(root.data,point) is True:
       return root

    #gen_data of root_cor and data_cor    
    if check == 0:
       root_cor = root.data.x
       data_cor = point.x
    elif check == 1:
       root_cor = root.data.y
       data_cor = point.y
    else:
       root_cor = root.data.z
       data_cor = point.z
    
    if root_cor > data_cor:
       return searchRec(root.left,point,level + 1)
    return searchRec(root.right,point,level + 1)
#==============================================================================

#==============================================================================
def minNode(root,left,right,level):
    res = root
    check = level % 3
    if check == 0:
       if left != None and left.data.x < res.data.x:
          res = left
       if right != None and right.data.x < res.data.x:
          res = right
    elif check == 1:
       if left != None and left.data.y < res.data.y:
          res = left
       if right != None and right.data.y < res.data.y:
          res = right
    else:
       if left != None and left.data.z < res.data.z:
          res = left
       if right != None and right.data.z < res.data.z:
          res = right
    return res
#==============================================================================

#==============================================================================
def maxNode(root,left,right,level):
    res = root
    check = level % 3
    if check == 0:
       if left != None and left.data.x > res.data.x:
          res = left
       if right != None and right.data.x > res.data.x:
          res = right
    elif check == 1:
       if left != None and left.data.y > res.data.y:
          res = left
       if right != None and right.data.y > res.data.y:
          res = right
    else:
       if left != None and left.data.z > res.data.z:
          res = left
       if right != None and right.data.z > res.data.z:
          res = right
    return res
#==============================================================================

#==============================================================================
#min point of level d:
def findMinRec(root,d,level):
    if root is None:
       return None
    check = level % 3
    if check == d:
       if root.left == None:
          return root
       return findMinRec(root.left,d,level + 1)    
    return minNode(root,findMinRec(root.left,d,level + 1),findMinRec(root.right,d,level + 1),d)
#==============================================================================

#==============================================================================
#min point of level d:
def findMaxRec(root,d,level):
    if root is None:
       return None
    check = level % 3
    if check == d:
       if root.right == None:
          return root
       return findMaxRec(root.right,d,level + 1)    
    return maxNode(root,findMaxRec(root.left,d,level + 1),findMaxRec(root.right,d,level + 1),d)
#==============================================================================

#==============================================================================
#bugs
def deleteRec(root,point,level):
    if root == None:
       return None
    check = level % 3

    #if current_node is point
    if checkPoint(root.data,point) is True:
       if root.right != None:
          min = findMinRec(root.right,check,level + 1)

          #copy point
          root.data.x = min.data.x
          root.data.y = min.data.y
          root.data.z = min.data.z

          root.right = deleteRec(root.right,min.data,level + 1)
       elif root.left != None:
          min = findMaxRec(root.left,check,level + 1)

          #copy point
          root.data.x = min.data.x
          root.data.y = min.data.y
          root.data.z = min.data.z

          root.right = deleteRec(root.left,min.data,level + 1)
          root.left = None
       else: #leaf node
          root = None
       return root

    #not contain point
    root_cor = 0.0
    data_cor = 0.0

    if check == 0:
       root_cor = root.data.x
       data_cor = point.x
    elif check == 1:
       root_cor = root.data.y
       data_cor = point.y
    else:
       root_cor = root.data.z
       data_cor = point.z
    
    if root_cor > data_cor:
       root.left = deleteRec(root.left,point,level + 1)
    else:
       root.right = deleteRec(root.right,point,level + 1)
    return root

#==============================================================================

#==============================================================================
def print2DUtil(root,space = 0) -> None:
    if root is None:
       return None
    space += 10
    if root.right is not None:
       print2DUtil(root.right,space)
    
    print("\n{}{},{},{}\n".format(' '*(space - 10),root.data.x,root.data.y,root.data.z))

    if root.left is not None:
       print2DUtil(root.left,space)
#==============================================================================

#==============================================================================
def clearRec(root) -> None:
    if root is None:
       return
    clearRec(root.left)
    clearRec(root.right)
    del root
#==============================================================================

#==============================================================================
def findNNRec(root,query_point,best_node,best_distance,depth):
    if not root:
       return best_distance
    dist = euclide_distance(root.data,query_point)

    if dist < best_distance:
       best_node.append(root)
       best_distance = dist
  
    check = depth % 3
    root_cor = 0.0
    data_cor = 0.0

    if check == 0:
       root_cor = root.data.x
       data_cor = query_point.x
    elif check == 1:
       root_cor = root.data.y
       data_cor = query_point.y
    else:
       root_cor = root.data.z
       data_cor = query_point.z

    if data_cor < root_cor:
       good_side = root.left
       bad_side = root.right
    else:
       good_side = root.right
       bad_side = root.left

    best_distance = findNNRec(good_side,query_point,best_node,best_distance,depth + 1)
    if abs(root_cor - data_cor) < best_distance:
       best_distance = findNNRec(bad_side,query_point,best_node,best_distance,depth + 1)
    return best_distance
    
#==============================================================================

"""class KdTree Definition"""

class KdTree:
      method = None
      def __init__(self):
          self.root = None
          self.size = 0
      
      def __del__(self):
          self.clear()

      def clear(self):
          clearRec(self.root)
          self.root = None
          self.size = 0
      
      def insert(self,data):
          self.root = insertRec(self.root,data,0)
          self.size += 1
      
      def search(self,point):
          node = searchRec(self.root,point,0)
          if node is None:
             return -1
          return node.data

      def findMin(self,d):
          return findMinRec(self.root,d,0)

      def delete(self,point):
          if self.search(point) == None:
             print("no element in Kd tree")
          else:
             self.root = deleteRec(self.root,point,0)
             self.size -= 1

      def buildKdTree(self,points):
          for point in points:
              self.insert(point)
      
      def buildKdTree_pre(self,points,begin,end):
          if begin >= end:
             return None
          mid = int((begin + end)/2)
          self.insert(points[mid])
          self.buildKdTree_pre(points,begin,mid)
          self.buildKdTree_pre(points,mid + 1,end) 

      def findkNN(self,query_point,k):
          stack = []
          findNNRec(self.root,query_point,stack,np.infty,0)
          for i in range(k):
              if len(stack) >= 1:
                 point = stack.pop().data
                 print("point{}: {},{},{} with euclide_distance: {}\n".format(i + 1,point.x,point.y,point.z,euclide_distance(query_point,point)))

      def getSize(self):
          return self.size
      
      def getRoot(self):
          return self.root
      
      def printTree(self):
          print2DUtil(self.root)
      
      def setMethod(self,method):
          self.method = method

"""#Kd Tree Definition (Split Value)
* class Node definition
* class KdTree-Split

class Node definition
"""

class Node:
      data = None #leaf_node <=> list else None
      median = None #median of the data in the node
      left = None  #left_child_node
      right = None #right_child_node
      parent = None #parent: parent_node  
      split = -1 #choose dimension to split


      def __init__(self, data = None, split = 0, median = None, left = None, right = None, parent = None):
          self.data = data
          self.split = split
          self.left = left
          self.right = right
          self.parrent = parent
          self.median = median

"""class KdTree-split"""

class KdTree_split:
      min_split = 1 #min_split: a leaf node contains no more than `min_split` data
      method = None
      def __init__(self, all_data, dimension, min_split = 1):
          self.min_split = min_split
          self.dimension = dimension
          self.size = len(all_data)
          self.root = self.buildKdTree(all_data)

      def buildKdTree(self,all_data):
          if len(all_data) == 0:
             return None
          
          if len(all_data) <= self.min_split:
             return Node(all_data,0,all_data[0],None,None)
          
          #dimension of variable
          split_value = self.get_variance_dimension(all_data)
          left_data,median,right_data = self.split(all_data,split_value)

          node = Node([],split_value,median,self.buildKdTree(left_data),self.buildKdTree(right_data))
          if node.left is None:
             node.left.parent = node
          if node.right is None:
             node.right.parent = node
          
          return node  

      def get_variance_dimension(self,all_data):
          means = [0.0] * self.dimension
          vars = [0.0] * self.dimension

          for data in all_data:
              for i in range(self.dimension):
                  means[i] += data[i]
          
          means = [mean / len(all_data) for mean in means]
          
          for data in all_data:
              for i in range(self.dimension):
                  vars[i] += (data[i] - means[i])**2
          
          max_var = -1
          value = -1
          for i,var in enumerate(vars):
              if var > max_var:
                 max_var = var
                 value = i

          return value

      def split(self,all_data,dim):
          length = len(all_data)
          
          #k is number of element in partial data
          def min_k(data,begin,end,k):
              #partition in quicksort
              pivot = data[begin]
              left = begin
              right = end
              while left < right:
                    while left < right and data[right][dim] > pivot[dim]:
                          right -= 1
                    
                    while left < right and data[left][dim] < pivot[dim]:
                          left += 1
                    if left >= right:
                       break
                    data[right],data[left] = data[left],data[right]
              data[left] = pivot
              return
              if left - begin + 1 <= k:
                 return 
              elif left - begin + 1 < k:
                 #choose last k - (left - begin + 1) element from partial data: left + 1 -> end
                 return min_k(data,left + 1,end,k - (left - begin + 1)) 
              else:
                 #choose k element from data with slice index form begin -> left -1
                 return min_k(data,begin,left - 1,k)
          mid = int(length/2)
          min_k(all_data,0,length - 1,mid)
          return all_data[0 : mid],all_data[mid - 1],all_data[mid:]

      def findkNN(self,data,k = 1):
          best_distance = np.inf
          dq = collections.deque()
          def find(root):
              nonlocal best_distance
              nonlocal dq
              if self.size == 0:
                 return None
              if root.left is None and root.right is None:  #leaf node
                 dist = euclide_distance_split(root.data,data)
                 if dist <= best_distance:
                    dq.append(root)
                    best_distance = dist
                    if len(dq) == k:
                       dq.popleft()
              else:
                 if data[root.split] <= root.median[root.split]:
                    good_side = root.left
                    bad_side = root.right
                 else:
                    good_side = root.right
                    bad_side = root.left
                 find(good_side)
                 if abs(data[root.split] - root.median[root.split]) <= best_distance:
                    find(bad_side)
          find(self.root)
          res = []
          for i in range(len(dq)):
              node_temp = dq.pop()
              res.append(node_temp)
              print("point{}: {} with euclide distance: {}".format(i + 1,np.squeeze(node_temp.data),euclide_distance_split(node_temp.data,data)))
        
      def search(self,data):
          self.findkNN(data)

      def getSize(self):
          return self.size
      
      def getRoot(self):
          return self.root
      
      def clear(self):
          def clearRec(root):
              if root is None:
                 return
              clearRec(root.left)
              clearRec(root.right)
              del root
              root = None
          clearRec(self.root)
          self.root = None
          self.size = 0

      def setMethod(self,method):
          self.method = method

"""#Data
* genData
* readData and convert Data
"""

#genData 3D-Data
def gen_data(n,filename):
    f = open(filename,"w")
    f.write("{} \n".format(n))
    for i in range(0,n):
        l_gen = []
        for j in range(0,3):
            g = round(random.uniform(0,100),3)
            l_gen.append(g)
        if i != n - 1:
            f.write("{} {} {}".format(l_gen[0],l_gen[1],l_gen[2]))
            f.write('\n')
        else:
            f.write("{} {} {}".format(l_gen[0],l_gen[1],l_gen[2]))
    f.close()

def readData(filename,method = "classic"):
    points = []
    with open(filename,'r',encoding='UTF-8') as f:
         n = int(f.readline())
         for i in range(n):
             line = f.readline()
             lst = re.split("\s+",line)
             if len(lst) >= 4:
                lst = lst[:-1]
             if method == "classic" or method == "re_balance":
                point = Point()
                point.set_x(round(float(lst[0]),3))
                point.set_y(round(float(lst[1]),3))
                point.set_z(round(float(lst[2]),3))
             else:
                point = [round(float(x),3) for x in lst]
             points.append(point)
    return points

"""#Main
**BUILD MODEL**
  *   build_method: gen_model with method and data
  *   3 - method:
      *   classic
      *   re_balance(preprocess data for balance tree)
      *   flip
 
**METHOD OF MODEL:**
  *   search
  *   findkNN
  *   getSize
  *   printTree (don't implement with split Kd-tree)




"""

#build_method
def build_method(filename,method = "classic"):
     points = readData(filename,method)
     if method == "re_balance":
        preprocess(points,0,0,len(points))
        model = KdTree()
        model.buildKdTree_pre(points,0,len(points))
     elif method == "classic":
        model = KdTree()
        model.buildKdTree(points)
     else: 
        model = KdTree_split(points,3)
     model.setMethod(method)
     return model,points

def search(query_point,model):
    method = model.method
    if method == "split":
       model.search(query_point)
    else:
       point_temp = model.search(query_point)
       if point_temp != -1:
          print("{},{},{}\n".format(point_temp.x,point_temp.y,point_temp.z))

def gen_query_point(method):
    l_gen = []
    for j in range(0,3):
        g = round(random.uniform(0,100),3)
        l_gen.append(g)
    if method == "classic" or method == "re_balance":
       point = Point()
       point.set_x(round(float(l_gen[0]),3))
       point.set_y(round(float(l_gen[1]),3))
       point.set_z(round(float(l_gen[2]),3))
    else:
       point = [round(float(x),3) for x in l_gen]
    return point

# filename = "a.txt"
# model,points = build_method(filename,"split")
# # search(points[0],model)
# point = gen_query_point("split")
# model.findkNN(point,3)