-
Notifications
You must be signed in to change notification settings - Fork 3
/
quadtree.py
executable file
·385 lines (283 loc) · 12.9 KB
/
quadtree.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
#!/usr/bin/env python
import utils
import geometry
from geometry import x,y
# import enum
def as_box( quadrant ):
""""Convert a quadrant of the form: ((x_min,y_min),width) to a box: ((x_min,y_min),(x_max,y_max))."""
width = quadrant[1]
minp = quadrant[0]
maxp = tuple(xy+width for xy in minp)
assert( x(minp) <= x(maxp) and y(minp) <= y(maxp) )
return (minp,maxp)
def as_rect( quadrant ):
""""Convert a quadrant of the form: ((x_min,y_min),width) to a rectangle: ((x0,y0),(x1,y1),(x2,y2),(x3,y3))."""
qx,qy = quadrant[0]
w = quadrant[1]
return [(qx,qy),(qx+w,qy),(qx+w,qy+w),(qx,qy+w)]
class QuadTree(object):
def __init__( self, points = [] ):
"""Build a quadtree on the given set of points.
Points must be an iterable containing 2-tuples of the form: (x,y)"""
# Initialize the root quadrant as the box around the points
self.root, self.quadrants = self.init( points = points )
# Each leaf of the quadtree may contains one resident point.
self.residents = { self.root: None }
# Each node of the quadtree may contains 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()
# Choose one of the two available functions for walking the tree:
# self.walk = self.recursive_walk
self.walk = self.iterative_walk
# Generate the complete tree.
self.build( points )
def init( self, quadrant = None, box = None, points = None ):
"""Initialize the root quadrant with the given quadrant, 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 a box, a quadrant or points")
# 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]
assert( x(minp) <= x(minp)+width and y(minp) <= y(minp)+width )
# There is always the root quadrant in the list of available ones.
root = (minp,width)
quadrants = [ root ]
return root,quadrants
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 = 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] )
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.
Move the existing resident to the correct child."""
# 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 origin in ( (qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy) ):
# Create a child quadrant of half its width
q = (origin, w)
self.quadrants.append(q)
# Default resident to None, because we will test for this key later on.
self.residents[q] = None
# Add this new child to the current parent.
self.children[quadrant].append(q)
# This new quadrant has no child.
self.children[q] = []
# 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 )
# FIXME use a recursive walk and prune branches with the Out status.
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, quadrant = None, depth = 0 ):
"""Return a string representing the quadtree in a JSON-like format."""
# Default to the root quadrant
if not quadrant:
quadrant = self.root
head = " "*depth
r = head+"{"
quadrep = '"origin" : %s, "width" : %f' % quadrant
if self.residents[quadrant]: # external
r += ' "resident" : %s, \t%s },\n' % (self.residents[quadrant],quadrep)
elif self.children[quadrant]: # internal
r += ' "children_ids" : %s, \t%s, "children" : [\n' % (self.children[quadrant],quadrep)
for child in self.children[quadrant]:
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 the set of points attached to the quadtree.
In a random order."""
return [p for p in self.residents.values() if p is not None]
def covers( self, this, that ):
"""Return true if the given quadrants does intersects each other."""
# Convert quadrants ((x,y),w) as box ((a,b),(c,d)).
this_box = as_box(this)
that_box = as_box(that)
# Convert boxes as list of edges.
this_segments = tuple(utils.tour(as_rect(this)))
that_segments = tuple(utils.tour(as_rect(that)))
# If at least one of the segment of "this" intersects with "that".
intersects = any( geometry.segment_intersection(s0,s1) for s0 in this_segments for s1 in that_segments )
# Transform nested list of segments in flat list of points without any duplicates.
this_points = as_rect(this)
that_points = as_rect(that)
# If all the points of "this" are inside "that".
# Note: what we would want to test here is if ALL the points are comprised,
# as the case where at least one is already tested by the "intersects" stage.
# But we use an "any" anyway, because it is sufficient in this case and
# that testing all the points takes more time.
this_in = any( geometry.in_box(p,this_box) for p in that_points )
that_in = any( geometry.in_box(p,that_box) for p in this_points )
return intersects or this_in or that_in
def query( self, query_quad, at_quad = None ):
"""Return all the points (currently attached to the quad tree) that are located within the query_quad quadrant."""
if not at_quad:
at_quad = self.root
query_box = as_box(query_quad)
# If we ask for a quadrant that intersects with the current one.
if self.covers( query_quad, at_quad ):
# If the current quadrant contains sub-quadrants.
if len(self.children[at_quad]) > 0:
# Then go explore them.
points = []
for quad in self.children[at_quad]:
points += self.query(query_quad,quad)
return points
else:
# Else, just return the point within the current quadrant.
resident = self.residents[at_quad]
if resident:
if geometry.in_box(resident,query_box):
# In a list, because we will concatenate.
return [resident]
# If there is no intersection, there is no points.
return []
# Pythonesque API:
def __getitem__( self, quadrant ):
"""Return all the points that are located within the given quadrant.
Ex.: points = quad[quad.root] # get all the points"""
return self.query(quadrant,self.root)
def __iter__(self):
"""Iterate over the attached points."""
return iter(self.points())
def __call__(self, points):
"""Append all the given points in the quadtree."""
self.build(points)
def __len__(self):
"""Return the number of points attached to the quad tree."""
return len(self.points())
def __repr__(self):
"""Return a string representing the quadtree in a JSON-like format."""
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 points in the quadtree / %i points\n" % (len(quad), len(points)) )
fig = plot.figure()
ax = fig.add_subplot(111)
ax.set_aspect('equal')
# Plot the whole quad tree and its points.
# Iterating over the quadtree will generate points, thus list(quad) is equivalent to quad.points()
uberplot.scatter_points( ax, list(quad), facecolor="green", edgecolor="None")
for q in quad.quadrants:
edges = list( utils.tour(as_rect(q)) )
uberplot.plot_segments( ax, edges, edgecolor = "blue", alpha = 0.1, linewidth = 2 )
# Plot a random query on the quad tree.
# Remember a quadrant is ( (orig_y,orig_y), width )
minp = ( round(random.uniform(-n,n),2), round(random.uniform(-n,n),2) )
rand_quad = ( minp, round(random.uniform(0,n),2) )
# Asking for a quadrant will query the quad tree and return the corresponding points.
uberplot.scatter_points( ax, quad[rand_quad], facecolor="None", edgecolor="red", alpha=0.5, linewidth = 2 )
edges = list( utils.tour(as_rect(rand_quad)) )
uberplot.plot_segments( ax, edges, edgecolor = "red", alpha = 0.5, linewidth = 2 )
plot.show()
assert(len(points) == len(quad))