# import enum import geometry from geometry import x,y def next_id( existing_ids ): i = 0 while i in existing_ids: i += 1 return i class QuadTree(object): def __init__( self, points = [] ): # Data structures to handle the quadtree # 0 is the root quadrant self.quadrants = { 0: None } self.widths = { 0: None } self.occupants = { 0: None } # Quadrants may have four children self.children = { 0: [] } # Initialize the root quadrant as the box around the points self.init(points = points) # Status of quadrants # class Status(enum.Enum): class Status: Leaf = 1 Node = 2 Empty = 3 Out = 4 self.Status = Status() self.walk = self.iterative_walk self.build(points) def init( self, quadrant = None, box = None, points = None ): if len([k for k in (box,points,quadrant) if k]) > 1: raise BaseException("ERROR: you should specify only one of the options") # Initialize the root quadrant as the box around the points if box: minp,maxp = box w = max( x(maxp)-x(minp), y(maxp)-y(minp) ) self.widths[0] = w self.quadrants[0] = minp elif points: minp,maxp = geometry.box( points ) w = max( x(maxp)-x(minp), y(maxp)-y(minp) ) self.widths[0] = w self.quadrants[0] = minp elif quadrant: self.quadrants[0] = quadrant[0] self.widths[0] = quadrant[1] else: raise BaseException("ERROR: you should specify a box, a quadrant or points") def appendable( self, p, quad ): """Return Status.Empty if the given point can be appended in the given quadrant.""" assert(p is not None) assert(len(p) == 2) minp = self.quadrants[quad] assert(minp is not None) assert(len(minp) == 2) w = self.widths[quad] maxp = tuple(i+w for i in minp) box = (minp,maxp) # if the point lies inside the given quadrant if geometry.in_box( p, box): if self.occupants[quad]: # external: a quadrant that already contains a point assert( not self.children[quad] ) # print("is external leaf") return self.Status.Leaf elif self.children[quad]: # internal: a quadrant that contains other quadrants # print("is internal node") return self.Status.Node else: # empty: there is not point yet in this quadrant # print("is empty") return self.Status.Empty else: # point is out of the quadrant # print("is out") return self.Status.Out def split(self, quadrant ): """Split an existing quadrant in four children quadrants. Spread existing occupants to the children.""" # We cannot split a quadrant if it already have sub-quadrants if quadrant != 0: assert( not self.children[quadrant] ) qx, qy = self.quadrants[quadrant] w = self.widths[quadrant] / 2 # For each four children quadrant's origins for orig in ((qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy)): # Create a child quadrant of half its width id = next_id( self.quadrants ) self.widths[id] = w self.quadrants[id] = orig self.occupants[id] = None # add a new child to the current parent self.children[quadrant].append(id) self.children[id] = [] # Move the occupant to the related children node p = self.occupants[quadrant] if p is not None: # Find the suitable children quadrant for child in self.children[quadrant]: if self.appendable(p,child) == self.Status.Empty: self.occupants[child] = p break # Forget we had occupant here # Do not pop the key, because we have tests on it elsewhere self.occupants[quadrant] = None assert( len(self.children[quadrant]) == 4 ) def append( self, point, quadrant = 0 ): """Try to inset the given point in the existing quadtree, under the given quadrant. The default quadrant is the root one (0). Returns True if the point was appended, False if it is impossible to append it.""" assert(quadrant in self.quadrants) # The point should not be out the root quadrant assert( self.appendable(point,0) != self.Status.Out ) for q in self.walk(quadrant): status = self.appendable( point, q ) if status == self.Status.Leaf: # Create sub-quadrants self.split(q) # Try to attach the point in children quadrants, recursively for child in self.children[q]: # print("consider children %i" % child) if self.append( point, child ): return True elif status == self.Status.Empty: # add the point as an occupant of the quadrant q self.occupants[q] = point # print("%s appended at %i" % (point,q)) return True return False def build(self, points): """append all the given points in the quadtree.""" for p in points: self.append(p) # assert( len(points) == len(self) ) def iterative_walk(self, at_quad = 0 ): # First, consider the root quadrant yield at_quad # then children of the root quadrant quads = list(self.children[at_quad]) for child in quads: yield child # then children of the current child quads.extend( self.children[child] ) def recursive_walk(self, at_quad = 0 ): yield at_quad for child in self.children[at_quad]: for q in self.recursive_walk(child): yield q def repr(self,quad=0,depth=0): head = " "*depth r = head+"{" quadrep = '"origin" : %s, "width" : %f' % (self.quadrants[quad],self.widths[quad]) if self.occupants[quad]: # external r += ' "id" : %i, "occupant" : %s, \t%s },\n' % (quad,self.occupants[quad],quadrep) elif self.children[quad]: # internal r += ' "id" : %i, "children_ids" : %s, \t%s, "children" : [\n' % (quad,self.children[quad],quadrep) for child in self.children[quad]: r += self.repr(child, depth+1) r+="%s]},\n" % head else: # empty r += ' "id" : %i, "occupant" : (), \t\t\t%s},\n' % (quad,quadrep) return r def points( self ): # return [self.occupants[q] for q in self.occupants if self.occupants[q] is not None] return [p for p in self.occupants.values() if p is not None] def __iter__(self): return iter(self.points()) def __call__(self, points): self.build(points) def __len__(self): """Return the number of points attached to the quad tree.""" return len(self.points()) def __repr__(self): return self.repr() if __name__ == "__main__": import sys import math import random import utils import uberplot import matplotlib.pyplot as plot if len(sys.argv) == 2: seed = sys.argv[1] else: seed = None random.seed(seed) n=200 points = [ ( round(random.uniform(-n,n),2),round(random.uniform(-n,n),2) ) for i in range(n) ] quad = QuadTree( points ) print(quad) sys.stderr.write( "%i attached points / %i appended points\n" % (len(quad), len(points)) ) fig = plot.figure() ax = fig.add_subplot(111) ax.set_aspect('equal') uberplot.scatter_points( ax, points, facecolor="red", edgecolor="red") uberplot.scatter_points( ax, quad.points(), facecolor="green", edgecolor="None") for q in quad.quadrants: qx, qy = quad.quadrants[q] w = quad.widths[q] box = [(qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy)] edges = list( utils.tour(box) ) uberplot.plot_segments( ax, edges, edgecolor = "blue", alpha = 0.1, linewidth = 2 ) plot.show() assert(len(points) == len(quad))