Reputation: 53
Ive created a simple quad tree class in python that initialises a tree, inserts a few bodies and attempts to plot them, alongside the boundaries it has made. This works for anything in the Quadrant 2 (South West), and its following subdivisions, however, Whenever I try to add bodies4 and bodies5 which should show up in Quadrant 3 (South East), nothing shows up. However, the node for quadrant 3 is created, but the bodies aren't added. Ive used recursive methods to repeatedly insert and subdivide, but that hasn't seemed to help:
import matplotlib.pyplot as plt
class Body:
def __init__(self, x, y):
self.x = x
self.y = y
self.isNode = False
class Node:
def __init__(self, boundx, boundy):
self.children = [None, None, None, None]
self.isNode = True
self.boundx = boundx
self.boundy = boundy
self.width = boundx * 2
self.height = boundy * 2
def insert(self, body, position):
self.children[position] = body
class Area:
def __init__(self, width, height):
self.width = width
self.height = height
class QuadTree:
def __init__(self, bounds):
self.size = 4
self.nodes = [None, None, None, None]
self.bounds = bounds
def insert(self, body):
target_quadrant = self.get_quadrant(body.x, body.y, self.bounds.width, self.bounds.height)
if self.nodes[target_quadrant] is None:
self.nodes[target_quadrant] = body
elif self.nodes[target_quadrant].isNode == False:
bodies = [self.nodes[target_quadrant], body]
self.nodes[target_quadrant] = self.subdivide(bodies, self.bounds, target_quadrant)
elif self.nodes[target_quadrant].isNode == True:
self.recursive_insert(body, self.nodes[target_quadrant])
def recursive_insert(self, body, node):
target_quadrant = self.get_quadrant(body.x, body.y, node.boundx, node.boundy)
if node.children[target_quadrant] is None:
node.children[target_quadrant] = body
elif node.children[target_quadrant].isNode == False:
bodies = [node.children[target_quadrant], body]
node.children[target_quadrant] = self.subdivide(bodies, node, target_quadrant)
elif node.children[target_quadrant].isNode == True:
self.recursive_insert(body, node.children[target_quadrant])
def get_offset(self, newBound, target_quadrant):
if target_quadrant == 0:
offset_x = newBound.width / 2
offset_y = newBound.height / 2
elif target_quadrant == 1:
offset_x = newBound.width / 2
offset_y = 0
elif target_quadrant == 2:
offset_x = 0
offset_y = 0
elif target_quadrant == 3:
offset_x = 0
offset_y = newBound.height / 2
return offset_x, offset_y
def subdivide(self, bodies, newBound, target_quadrant):
offset_x, offset_y = self.get_offset(newBound, target_quadrant)
new_Node = Node(newBound.width / 2, newBound.height / 2)
target_quadrants = []
for body in bodies:
target_quadrant = self.get_quadrant(body.x, body.y, newBound.width / 2, newBound.height / 2)
if target_quadrant == "Out of bounds":
continue
if target_quadrant in target_quadrants:
sub_quad_bodies = [b for b in bodies if self.get_quadrant(b.x, b.y, newBound.width / 2, newBound.height / 2) == target_quadrant]
sub_bound = Area(newBound.width / 2, newBound.height / 2)
new_Node.children[target_quadrant] = self.subdivide(sub_quad_bodies, sub_bound, target_quadrant)
else:
new_Node.insert(body, target_quadrant)
target_quadrants.append(target_quadrant)
return new_Node
def get_quadrant(self, x, y, width, height):
if x < 0 or y < 0 or x >= width or y >= height:
return "Out of bounds"
elif x < width / 2:
if y < height / 2:
return 2
else:
return 0
else:
if y < height / 2:
return 3
else:
return 1
def plot_bodies(self):
plt.figure(figsize=(8, 8))
plt.xlim(0, self.bounds.width)
plt.ylim(0, self.bounds.height)
self.plot_quadrants(self.bounds, self.nodes)
for node in self.nodes:
if node is not None:
self.plot_bodies_recursive(node)
plt.xlabel('X')
plt.ylabel('Y')
plt.title('QuadTree Bodies')
plt.show()
def plot_quadrants(self, bounds, nodes):
if bounds.width == 1 or bounds.height == 1:
return
else:
plt.plot([bounds.width / 2, bounds.width / 2], [0, bounds.height], color='gray', linestyle='--')
plt.plot([0, bounds.width], [bounds.height / 2, bounds.height / 2], color='gray', linestyle='--')
for node in nodes:
if node is not None and isinstance(node, Node):
self.plot_quadrants(Area(bounds.width / 2, bounds.height / 2), node.children)
def plot_bodies_recursive(self, node):
if node is not None:
if not node.isNode:
plt.scatter(node.x, node.y, color='red')
else:
for child in node.children:
self.plot_bodies_recursive(child)
if __name__ == "__main__":
bound = Area(128, 128)
body = Body(1, 1)
body2 = Body(5, 6)
body3 = Body(2, 4)
body4 = Body(80, 12)
body5 = Body(80, 60)
tree = QuadTree(bound)
tree.insert(body)
tree.insert(body2)
tree.insert(body3)
tree.insert(body4)
tree.insert(body5)
tree.plot_bodies()
Upvotes: 0
Views: 21