A quadtree is a tree data structure in which each node has exactly four children. It is a particularly efficient way to store elements when you need to quickly find them according to their x-y coordinates.
A common problem with elements in quadtrees is to detect pairs of elements which are closer than a definite threshold.
The proposed implementation efficiently addresses this problem.
In [1]:
from smartquadtree import Quadtree
As you instantiate your quadtree, you must specify the center of your space then the height and width.
In [2]:
q = Quadtree(0, 0, 10, 10)
The output of a quadtree on the console is pretty explicit. (You can refer to next section for the meaning of "No mask set")
In [3]:
q
Out[3]:
You can easily insert elements from which you can naturally infer x-y coordinates (e.g. tuples or lists)
In [4]:
q.insert((1, 2))
q.insert((-3, 4))
q
Out[4]:
No error is raised if the element you are trying to insert is outside the scope of the quadtree. But it won't be stored anyway!
In [5]:
q.insert((-20, 0))
q
Out[5]:
If you want to insert other Python objects, be sure to provide get_x() and get_y() methods to your class!
In [6]:
class Point(object):
def __init__(self, x, y, color):
self.x = x
self.y = y
self.color = color
def __repr__(self):
return "(%.2f, %.2f) %s" % (self.x, self.y, self.color)
def get_x(self):
return self.x
def get_y(self):
return self.y
You cannot insert elements of a different type from the first element inserted.
In [7]:
q.insert(Point(2, -7, "red"))
But feel free to create a new one and play with it:
In [8]:
point_quadtree = Quadtree(5, 5, 5, 5)
point_quadtree.insert(Point(2, 7, "red"))
point_quadtree
Out[8]:
In [9]:
from random import random
q = Quadtree(0, 0, 10, 10, 16)
for a in range(50):
q.insert([random()*20-10, random()*20-10])
The print function does not display all elements and uses the __repr__() method of each element.
In [10]:
print(q)
We can write our own iterator and print each element we encounter the way we like.
In [11]:
from __future__ import print_function
for p in q.elements():
print ("[%.2f, %.2f]" % (p[0], p[1]), end=" ")
It is easy to filter the iteration process and apply the function only on elements inside a given polygon. Use the set_mask() method and pass a list of x-y coordinates. The polygon will be automatically closed.
In [12]:
q.set_mask([(-3, -7), (-3, 7), (3, 7), (3, -7)])
print(q)
The same approach can be used to count the number of elements inside the quadtree.
In [13]:
print (sum (1 for x in q.elements()))
print (sum (1 for x in q.elements(ignore_mask=True)))
As a mask is set on the quadtree, we only counted the elements inside the mask. You can use the size() method to count elements and ignore the mask by default. Disabling the mask with set_mask(None) is also a possibility.
In [14]:
print ("%d elements (size method)" % q.size())
print ("%d elements (don't ignore the mask)" % q.size(False))
q.set_mask(None)
print ("%d elements (disable the mask)" % q.size())
In [15]:
%matplotlib inline
from matplotlib import pyplot as plt
q = Quadtree(5, 5, 5, 5, 10)
for a in range(200):
q.insert([random()*10, random()*10])
fig = plt.figure()
plt.axis([0, 10, 0, 10])
q.set_mask(None)
for p in q.elements():
plt.plot([p[0]], [p[1]], 'o', color='lightgrey')
q.set_mask([(3, 3), (3, 7), (7, 7), (7, 3)])
for p in q.elements():
plt.plot([p[0]], [p[1]], 'ro')
_ = plt.plot([3, 3, 7, 7, 3], [3, 7, 7, 3, 3], 'r')
Iterating on pairs of neighbouring elements is possible through the neighbour_elements() function. It works as a generator and yields pair of elements, the first one being inside the mask (if specified), the second one being in the same cell or in any neighbouring cell, also in the mask.
Note that if (a, b) is yielded by neighbour_elements(), (b, a) will be omitted from future yields.
In [16]:
q = Quadtree(5, 5, 5, 5, 10)
q.set_limitation(2) # do not create a new subdivision if one side of the cell is below 2
for a in range(200):
q.insert([random()*10, random()*10])
fig = plt.figure()
plt.axis([0, 10, 0, 10])
for p in q.elements():
plt.plot([p[0]], [p[1]], 'o', color='lightgrey')
q.set_mask([(1, 1), (4, 1), (5, 4), (2, 5), (1, 1)])
for p in q.elements():
plt.plot([p[0]], [p[1]], 'o', color='green')
for p1, p2 in q.neighbour_elements():
if ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2 < 1):
plt.plot([p1[0]], [p1[1]], 'o', color='red')
plt.plot([p2[0]], [p2[1]], 'o', color='red')
plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'red')
_ = plt.plot([1, 4, 5, 2, 1], [1, 1, 4, 5, 1], 'r')
In [ ]: