Pella86

kd-tree beginner

Mar 6th, 2012
166
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 4.19 KB | None | 0 0
  1. #implementation of a kd tree... or at least a try
  2.  
  3. #based on the tutorial of "Range Searching Using Kd Tree, Hemant M. Kakade, 25 - 08 - 2005"
  4.  
  5. class V2:
  6. #general point class
  7.     def __init__(self, x = 0, y = 0):
  8.         self.x = x
  9.         self.y = y
  10.    
  11.     def __str__(self):
  12.         s = "("+str(self.x) + ":" + str(self.y)+")"
  13.         return s
  14.  
  15. class Node:
  16. #node class to store the points
  17.     def __init__(self, v2 = V2()):
  18.         self.data = v2
  19.         self.left = None
  20.         self.right = None
  21.  
  22.  
  23. class kdTree:
  24.    
  25.     def __init__(self):
  26.         root = Node()
  27.    
  28.     def is_leaf(self, node):
  29.         return node.left == None and node.right == None    
  30.  
  31. #The median point function runs through all the data_set   
  32.     def find_median_point_x(self,point_set):
  33.         #input: a list of point
  34.         #output: the point that has the x median value
  35.         # sort the point set by x, then pick the point lying in the middle
  36.         # of the length
  37.         sorted_set = sorted(point_set,key = lambda point: point.x)
  38.         median_index = int(len(sorted_set)/2.)
  39.         if median_index == 1: median_index = 0
  40.         return sorted_set[median_index]
  41.  
  42.     def find_median_point_y(self,point_set):
  43.         #see find_median_point_x
  44.         sorted_set = sorted(point_set,key = lambda point: point.y)
  45.         median_index = int(len(sorted_set)/2.)
  46.         if median_index == 1: median_index = 0
  47.         return sorted_set[median_index]
  48.    
  49.    
  50.     def separate_vertical(self,point_set, median):
  51.         #input: list of point and a point that is the median
  52.         #output: two list containing points who belong to right or left
  53.         #given a vertical line, traced  by a point, that separate the 2d space,
  54.         #sort in two different "buckets" the points which lies to the left or to the right of
  55.         #that given line
  56.         left_sub = []
  57.         right_sub = []
  58.         for point in point_set:
  59.             if point.x <= median.x:
  60.                 left_sub.append(point)
  61.             else:
  62.                 right_sub.append(point)
  63.         return (left_sub,right_sub)
  64.  
  65.     def separate_horizontal(self,point_set, median):
  66.         #see separate vertical, but finds ab
  67.         left_sub = []
  68.         right_sub = []
  69.         for point in point_set:
  70.             if point.y <= median.y:
  71.                 left_sub.append(point)
  72.             else:
  73.                 right_sub.append(point)
  74.         return (left_sub,right_sub)
  75.        
  76. #build the tree given a point_set  
  77.     def build(self, point_set, depth):
  78.         P1 = []
  79.         P2 = []
  80.         median = None
  81.         #stop the recursion if the data set length is 1
  82.         #creating a leaf containing a point
  83.         if len(point_set) == 1:
  84.             node = Node(point_set[0])
  85.             return node
  86.         elif depth % 2 == 0: #if depth is even
  87.             median = self.find_median_point_x(point_set) #find the median vertical line
  88.             subsets = self.separate_vertical(point_set, median) #divide the point left or right
  89.             P1 = subsets[0]
  90.             P2 = subsets[1]
  91.         else:
  92.             median = self.find_median_point_y(point_set)
  93.             subsets = self.separate_horizontal(point_set, median)
  94.             P1 = subsets[0]
  95.             P2 = subsets[1]
  96.         #recurse to create more nodes
  97.         NodeL = self.build(P1,depth + 1)
  98.         NodeR = self.build(P2,depth + 1)
  99.        
  100.         #build the split node by storing the point that splitted the space there
  101.         #and append the two other nodes
  102.         node = Node()
  103.         node.data = median
  104.         node.left = NodeL
  105.         node.right = NodeR
  106.         return node
  107.  
  108. #Auxilliary functions  
  109.     def report_subtree(self, node, points):
  110.         #build a list of points starting from a node
  111.         if node.left != None:
  112.             if self.is_leaf(node.left):
  113.                 n1 = node.left
  114.                 points.append(n1.data)
  115.             else:
  116.                 self.report_subtree(node.left, points)
  117.         if node.right != None:
  118.             if self.is_leaf(node.right):
  119.                 n1 = node.right
  120.                 points.append(n1.data)
  121.             else:
  122.                 self.report_subtree(node.right, points)
  123.         return points
  124.  
  125.     def print_node(self, node , depth, s):
  126.         #print the subtree
  127.         s = depth*" " + str(node.data) + "\n"
  128.         if node.left != None:
  129.             s += self.print_node(node.left, depth + 1, s)
  130.         if node.right != None:
  131.             s += self.print_node(node.right, depth + 1, s)
  132.         return s
  133.            
  134.     def print_tree(self):
  135.         print self.print_node(self.root, 0, "")
  136.    
  137.  
  138.                
  139.            
  140.  
  141. if __name__ == "__main__":
  142.     data_set = []
  143.     for i in range(100):
  144.         data_set.append(V2(i,i+2))
  145.  
  146.     kdtree = kdTree()
  147.     kdtree.root = kdtree.build(data_set, 0)
  148.     kdtree.print_tree()
  149.  
  150.     points = []
  151.     points = kdtree.report_subtree(kdtree.root.left, points)
  152.     for point in points:
  153.         print point,
  154.    
  155.     print "Hello world!"
Advertisement
Add Comment
Please, Sign In to add comment