292 lines
8.6 KiB
Python
292 lines
8.6 KiB
Python
#!/usr/bin/env python
|
|
|
|
import geometry
|
|
from geometry import x,y
|
|
|
|
# import enum
|
|
|
|
class QuadTree(object):
|
|
|
|
def __init__( self, points = [] ):
|
|
"""Build a quadtree on the given set of points."""
|
|
|
|
# Initialize the root quadrant as the box around the points
|
|
self.init( points = points )
|
|
|
|
# Data structures to handle the quadtree
|
|
self.residents = { self.root: None }
|
|
|
|
# Quadrants may have four children
|
|
self.children = { self.root: [] }
|
|
|
|
# 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 ):
|
|
"""Initialize the root quadrant with the given quadrant ((x,y),width), the given box or the given set of points."""
|
|
|
|
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 given box
|
|
if box:
|
|
minp,maxp = box
|
|
width = max( x(maxp)-x(minp), y(maxp)-y(minp) )
|
|
|
|
# Initialize the root quadrant as the box around the points
|
|
elif points:
|
|
minp,maxp = geometry.box( points )
|
|
width = max( x(maxp)-x(minp), y(maxp)-y(minp) )
|
|
|
|
# Initialize the root quadrant as the given origin point and width
|
|
elif quadrant:
|
|
minp = quadrant[0]
|
|
width = quadrant[1]
|
|
|
|
else:
|
|
raise BaseException("ERROR: you should specify a box, a quadrant or points")
|
|
|
|
# There is always the root quadrant in the list of available ones.
|
|
self.root = (minp,width)
|
|
self.quadrants = [ self.root ]
|
|
|
|
|
|
def as_box( self, quadrant );
|
|
width = quadrant[1]
|
|
maxp = tuple(xy+width for xy in quadrant[0])
|
|
return (quadrant[0],maxp)
|
|
|
|
|
|
def status( self, point, quadrant ):
|
|
"""Return Status.Empty if the given point can be appended in the given quadrant."""
|
|
|
|
assert(point is not None)
|
|
assert(len(point) == 2)
|
|
assert(quadrant is not None)
|
|
assert(len(quadrant) == 2)
|
|
|
|
box = self.as_box( quadrant )
|
|
|
|
# if the point lies inside the given quadrant
|
|
if geometry.in_box( point, box):
|
|
if self.residents[quadrant]:
|
|
# external: a quadrant that already contains a point
|
|
assert( not self.children[quadrant] )
|
|
# print("is external leaf")
|
|
return self.Status.Leaf
|
|
elif self.children[quadrant]:
|
|
# internal: a quadrant that contains other quadrants
|
|
return self.Status.Node
|
|
else:
|
|
# empty: there is not point yet in this quadrant
|
|
return self.Status.Empty
|
|
else:
|
|
# point is out of the quadrant
|
|
return self.Status.Out
|
|
|
|
|
|
def split(self, quadrant ):
|
|
"""Split an existing quadrant in four children quadrants.
|
|
|
|
Spread existing residents to the children."""
|
|
|
|
# We cannot split a quadrant if it already have sub-quadrants
|
|
if quadrant != self.root:
|
|
assert( not self.children[quadrant] )
|
|
|
|
qx, qy = quadrant[0]
|
|
w = quadrant[1] / 2
|
|
|
|
# For each four children quadrant's origins
|
|
self.children[quadrant] = []
|
|
for orig in ((qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy)):
|
|
q = (orig,w)
|
|
# Create a child quadrant of half its width
|
|
self.quadrants.append(q)
|
|
self.residents[q] = None
|
|
|
|
# Add a new child to the current parent.
|
|
self.children[quadrant].append(q)
|
|
# The new quadrant has no child.
|
|
self.children[q] = []
|
|
|
|
assert( len(self.children[quadrant]) == 4 )
|
|
|
|
# Move the resident to the related children node
|
|
p = self.residents[quadrant]
|
|
if p is not None:
|
|
# Find the suitable children quadrant
|
|
for child in self.children[quadrant]:
|
|
if self.status(p,child) == self.Status.Empty:
|
|
self.residents[child] = p
|
|
break
|
|
# Forget we had resident here
|
|
# Do not pop the key, because we have tests on it elsewhere
|
|
self.residents[quadrant] = None
|
|
|
|
|
|
|
|
def append( self, point, quadrant = None ):
|
|
"""Try to inset the given point in the existing quadtree, under the given quadrant.
|
|
|
|
The default quadrant is the root one.
|
|
Returns True if the point was appended, False if it is impossible to append it."""
|
|
|
|
# Default to the root quadrant
|
|
if not quadrant:
|
|
quadrant = self.root
|
|
assert(quadrant in self.quadrants)
|
|
|
|
# The point should not be out of the root quadrant
|
|
assert( self.status(point,self.root) != self.Status.Out )
|
|
|
|
for q in self.walk(quadrant):
|
|
status = self.status( 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]:
|
|
if self.append( point, child ):
|
|
return True
|
|
elif status == self.Status.Empty:
|
|
# add the point as an resident of the quadrant q
|
|
self.residents[q] = point
|
|
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 = None ):
|
|
|
|
# Default to the root quadrant
|
|
if not at_quad:
|
|
at_quad = self.root
|
|
|
|
# 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 = None ):
|
|
|
|
# Default to the root quadrant
|
|
if not at_quad:
|
|
at_quad = self.root
|
|
|
|
yield at_quad
|
|
for child in self.children[at_quad]:
|
|
for q in self.recursive_walk(child):
|
|
yield q
|
|
|
|
|
|
def repr(self, quad=None, depth=0):
|
|
|
|
# Default to the root quadrant
|
|
if not quad:
|
|
quad = self.root
|
|
|
|
head = " "*depth
|
|
r = head+"{"
|
|
quadrep = '"origin" : %s, "width" : %f' % quad
|
|
if self.residents[quad]: # external
|
|
r += ' "resident" : %s, \t%s },\n' % (self.residents[quad],quadrep)
|
|
elif self.children[quad]: # internal
|
|
r += ' "children_ids" : %s, \t%s, "children" : [\n' % (self.children[quad],quadrep)
|
|
for child in self.children[quad]:
|
|
r += self.repr(child, depth+1)
|
|
r+="%s]},\n" % head
|
|
else: # empty
|
|
r += ' "resident" : (), \t\t\t%s},\n' % (quadrep)
|
|
return r
|
|
|
|
|
|
def points( self ):
|
|
return [p for p in self.residents.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=20
|
|
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 points in the quadtree / %i 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")
|
|
uberplot.scatter_points( ax, list(quad), facecolor="green", edgecolor="None")
|
|
|
|
for q in quad.quadrants:
|
|
qx, qy = q[0]
|
|
w = q[1]
|
|
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))
|
|
|