Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #implementation of a kd tree... or at least a try
- #based on the tutorial of "Range Searching Using Kd Tree, Hemant M. Kakade, 25 - 08 - 2005"
- class V2:
- #general point class
- def __init__(self, x = 0, y = 0):
- self.x = x
- self.y = y
- def __str__(self):
- s = "("+str(self.x) + ":" + str(self.y)+")"
- return s
- class Node:
- #node class to store the points
- def __init__(self, v2 = V2()):
- self.data = v2
- self.left = None
- self.right = None
- class kdTree:
- def __init__(self):
- root = Node()
- def is_leaf(self, node):
- return node.left == None and node.right == None
- #The median point function runs through all the data_set
- def find_median_point_x(self,point_set):
- #input: a list of point
- #output: the point that has the x median value
- # sort the point set by x, then pick the point lying in the middle
- # of the length
- sorted_set = sorted(point_set,key = lambda point: point.x)
- median_index = int(len(sorted_set)/2.)
- if median_index == 1: median_index = 0
- return sorted_set[median_index]
- def find_median_point_y(self,point_set):
- #see find_median_point_x
- sorted_set = sorted(point_set,key = lambda point: point.y)
- median_index = int(len(sorted_set)/2.)
- if median_index == 1: median_index = 0
- return sorted_set[median_index]
- def separate_vertical(self,point_set, median):
- #input: list of point and a point that is the median
- #output: two list containing points who belong to right or left
- #given a vertical line, traced by a point, that separate the 2d space,
- #sort in two different "buckets" the points which lies to the left or to the right of
- #that given line
- left_sub = []
- right_sub = []
- for point in point_set:
- if point.x <= median.x:
- left_sub.append(point)
- else:
- right_sub.append(point)
- return (left_sub,right_sub)
- def separate_horizontal(self,point_set, median):
- #see separate vertical, but finds ab
- left_sub = []
- right_sub = []
- for point in point_set:
- if point.y <= median.y:
- left_sub.append(point)
- else:
- right_sub.append(point)
- return (left_sub,right_sub)
- #build the tree given a point_set
- def build(self, point_set, depth):
- P1 = []
- P2 = []
- median = None
- #stop the recursion if the data set length is 1
- #creating a leaf containing a point
- if len(point_set) == 1:
- node = Node(point_set[0])
- return node
- elif depth % 2 == 0: #if depth is even
- median = self.find_median_point_x(point_set) #find the median vertical line
- subsets = self.separate_vertical(point_set, median) #divide the point left or right
- P1 = subsets[0]
- P2 = subsets[1]
- else:
- median = self.find_median_point_y(point_set)
- subsets = self.separate_horizontal(point_set, median)
- P1 = subsets[0]
- P2 = subsets[1]
- #recurse to create more nodes
- NodeL = self.build(P1,depth + 1)
- NodeR = self.build(P2,depth + 1)
- #build the split node by storing the point that splitted the space there
- #and append the two other nodes
- node = Node()
- node.data = median
- node.left = NodeL
- node.right = NodeR
- return node
- #Auxilliary functions
- def report_subtree(self, node, points):
- #build a list of points starting from a node
- if node.left != None:
- if self.is_leaf(node.left):
- n1 = node.left
- points.append(n1.data)
- else:
- self.report_subtree(node.left, points)
- if node.right != None:
- if self.is_leaf(node.right):
- n1 = node.right
- points.append(n1.data)
- else:
- self.report_subtree(node.right, points)
- return points
- def print_node(self, node , depth, s):
- #print the subtree
- s = depth*" " + str(node.data) + "\n"
- if node.left != None:
- s += self.print_node(node.left, depth + 1, s)
- if node.right != None:
- s += self.print_node(node.right, depth + 1, s)
- return s
- def print_tree(self):
- print self.print_node(self.root, 0, "")
- if __name__ == "__main__":
- data_set = []
- for i in range(100):
- data_set.append(V2(i,i+2))
- kdtree = kdTree()
- kdtree.root = kdtree.build(data_set, 0)
- kdtree.print_tree()
- points = []
- points = kdtree.report_subtree(kdtree.root.left, points)
- for point in points:
- print point,
- print "Hello world!"
Advertisement
Add Comment
Please, Sign In to add comment