I'm trying to create a Quadtree structure in python for detecting collisions amongst polygons, and I've gotten pretty far (see end of the post). However, I realized that this structure only works for points, because I'm deciding which leaf to put the object in based off of its center.
So I need to figure out how to modify this quadtree so I can detect collisions of regions (like a circle!).
There's a few different ways I can think about doing this:
The option #2 (placing the object in multiple nodes), I feel like it'd be best to draw a bounding rectangle around the object, and then use the extents of the object to decide which leaf to put it in -- but this would require effectively inserting each object 4 times (one for each corner), and that seems like an inefficient way to approach the problem.
Any better suggestions?
class Quadtree(object):
"""
A simple Quadtree implementation. This works with any child object that has a getPosition() method that returns a list or tuple.
"""
def __init__(self, depth, min_x, min_y, max_x, max_y):
self.mDepth = depth
self.mMaxDepth = 4
self.mMinX = min_x
self.mMaxX = max_x
self.mMidX = (max_x - min_x) / 2.0
self.mMinY = min_y
self.mMaxY = max_y
self.mMidY = (max_y - min_y) / 2.0
self.mHasChildren = False
self.mMaxChildren = 8
self.mChildren = []
self.mSubtrees = []
def insert(self, newChild):
"""
Insert an object into the tree. Returns True if the insert was successful, False otherwise.
"""
if self.mSubtrees:
#if subtrees exist, add the child to the subtrees
index = getIndex(newChild)
if index != -1:
self.mSubtrees[index].insert(newChild)
return True
#if no subtrees exist, add the child to the child list.
self.mChildren.append(newChild)
#and then check if we need to split the tree
#if there are more than the max children, and we haven't maxed out the tree depth, and there are no subtrees yet
if len(self.mChildren) > self.mMaxChildren and self.mDepth < self.mMaxDepth and not self.mSubtrees:
split()
for child in self.mChildren:
index = getIndex(child)
if index != -1:
self.mSubtrees[index].insert(child)
return True
return False
def retrieveNeighbors(self, targetChild):
index = getIndex(targetChild)
if index != -1 and self.mSubtrees:
return self.mSubtrees[index].retrieve(targetChild)
return self.mChildren
def getIndex(self, child):
"""
Returns the index of the node that the object belongs to. Returns -1 if the object does not exist in the tree.
"""
index = -1
childPosition = child.getPosition()
#check if it fits in the top or bottom
isInTopQuadrant = childPosition[1] > self.mMidY and childPositoin[1] < self.mMaxY
#check if it fits left
if childPosition[0] < self.mMidX and childPosition[0] > self.mMinX:
if isInTopQuadrant:
index = 1
else:
index = 2
#check if it fits right
if childPosition[0] > self.mMidX and childPosition[0] < self.mMidX:
if isInTopQuadrant:
index = 0
else:
index = 3
return index
def split(self):
"""
Split the trees into four subtrees.
"""
#top right
self.mSubtrees.append(Quadtree(depth + 1, self.mMidX, self.mMidY, self.mMaxX, self.mMaxY))
#top left
self.mSubtrees.append(Quadtree(depth + 1, self.mMinX, self.mMidY, self.mMidX, self.mMaxY))
#bottom left
self.mSubtrees.append(Quadtree(depth + 1, self.mMinX, self.mMinY, self.mMidX, self.mMidY))
#bottom right
self.mSubtrees.append(Quadtree(depth + 1, self.mMidX, self.mMinY, self.mMaxX, self.mMidY))
def clear(self):
"""
Clears the quadtree of children, and all subtrees of children.
"""
self.mChildren[:] = []
self.mHasChildren = False
for tree in range(0,4):
if self.mSubtrees[tree].mHasChildren:
self.mSubtrees[tree].clear()
The best way I've found so far is to modify _get_index() to check the entire size of the particle -- if it doesn't entirely fit in a subtree, then it becomes a child of the parent node:
def _get_index(self, child):
"""
Returns the index of the node that the object belongs to. Returns -1 if the object does not exist in the tree.
"""
index = -1
child_position = child.get_position()
child_radius = child.get_radius()
if len(child_position) != 2:
print "Quadtree is only designed to handle two-dimensional objects! Object will not be added."
return index
#check if it fits in the top or bottom
is_in_top_quadrant = child_position[1] - child_radius > self.mid_y and child_position[1] + child_radius < self.max_y
#check if it fits left
if child_position[0] + child_radius < self.mid_x and child_position[0] - child_radius > self.min_x:
if is_in_top_quadrant:
index = 1
else:
index = 2
#check if it fits right
if child_position[0] - child_radius > self.mid_x and child_position[0] + child_radius < self.mid_x:
if is_in_top_quadrant:
index = 0
else:
index = 3
return index
Then retrieve_neighbors() can traverse the tree and add the children of each node it passes, all the way down to the leaf node.