159 lines
4.9 KiB
Python
159 lines
4.9 KiB
Python
from math import floor, ceil
|
|
|
|
|
|
# A utility class, a rectangle acting as boundaries of the quad tree
|
|
class Bounds:
|
|
def __init__(self, left, right, bottom, top):
|
|
if top < bottom:
|
|
top, bottom = bottom, top
|
|
|
|
if right < left:
|
|
right, left = left, right
|
|
|
|
self.left = left
|
|
self.right = right
|
|
self.bottom = bottom
|
|
self.top = top
|
|
self.width = right - left
|
|
self.height = top - bottom
|
|
|
|
def getWidth(self):
|
|
return self.width
|
|
|
|
def getHeight(self):
|
|
return self.height
|
|
|
|
def getLeft(self):
|
|
return self.left
|
|
|
|
def getRight(self):
|
|
return self.right
|
|
|
|
def getBottom(self):
|
|
return self.bottom
|
|
|
|
def getTop(self):
|
|
return self.top
|
|
|
|
|
|
# A QuadTree implemented to specifically handle panel Nodes
|
|
class NodeQuadTree():
|
|
MAX_OBJECTS = 100
|
|
MAX_LEVELS = 5
|
|
|
|
def __init__(self, level, bounds, variance):
|
|
if level < 1:
|
|
level = 1
|
|
|
|
self.level = level
|
|
self.variance = variance
|
|
|
|
self.bounds = bounds
|
|
if level == 1:
|
|
self.bounds = Bounds(floor(bounds.getLeft() - self.variance),
|
|
ceil(bounds.getRight() + self.variance),
|
|
floor(bounds.getBottom() - self.variance),
|
|
ceil(bounds.getTop() + self.variance))
|
|
|
|
self.nodeList = []
|
|
self.quads = [None, None, None, None]
|
|
|
|
def clear(self):
|
|
self.nodeList = []
|
|
for quad in self.quads:
|
|
if quad is not None:
|
|
quad.clear()
|
|
self.quads = [None, None, None, None]
|
|
|
|
def pointInside(self, point):
|
|
if point.x - self.variance < self.bounds.getLeft():
|
|
return False
|
|
if point.x + self.variance > self.bounds.getRight():
|
|
return False
|
|
if point.y - self.variance < self.bounds.getBottom():
|
|
return False
|
|
if point.y + self.variance > self.bounds.getTop():
|
|
return False
|
|
return True
|
|
|
|
def split(self):
|
|
left = self.bounds.getLeft()
|
|
right = self.bounds.getRight()
|
|
midLR = left + self.bounds.getWidth() / 2
|
|
|
|
bottom = self.bounds.getBottom()
|
|
top = self.bounds.getTop()
|
|
midBT = bottom + self.bounds.getHeight() / 2
|
|
|
|
self.quads[0] = NodeQuadTree(self.level + 1, Bounds(left, midLR, bottom, midBT), self.variance)
|
|
self.quads[1] = NodeQuadTree(self.level + 1, Bounds(midLR, right, bottom, midBT), self.variance)
|
|
self.quads[2] = NodeQuadTree(self.level + 1, Bounds(left, midLR, midBT, top), self.variance)
|
|
self.quads[3] = NodeQuadTree(self.level + 1, Bounds(midLR, right, midBT, top), self.variance)
|
|
|
|
# Returns which child index the point belongs in, or -1 if it doesn't fit completely within any
|
|
def getIndex(self, point):
|
|
for i in range(4):
|
|
if self.quads[i] is not None:
|
|
if self.quads[i].pointInside(point):
|
|
return i
|
|
return -1
|
|
|
|
# insert the node into the QuadTree, or one of its children
|
|
def insert(self, node):
|
|
# add it to our children, if they exist and the point fits
|
|
if self.quads[0] is not None:
|
|
index = self.getIndex(node.coordinate)
|
|
|
|
if index != -1:
|
|
self.quads[index].insert(node)
|
|
return
|
|
|
|
# else, add it to self
|
|
self.nodeList.append(node)
|
|
|
|
# too big? split into quads
|
|
if (len(self.nodeList) > NodeQuadTree.MAX_OBJECTS) and \
|
|
(self.level < NodeQuadTree.MAX_LEVELS) and \
|
|
(self.quads[0] is None):
|
|
|
|
self.split()
|
|
|
|
toKeep = []
|
|
for n in self.nodeList:
|
|
index = self.getIndex(n.coordinate)
|
|
|
|
if index != -1:
|
|
self.quads[index].insert(n)
|
|
else:
|
|
toKeep.append(n)
|
|
|
|
self.nodeList = toKeep
|
|
|
|
# Return a list of all possible nodes that can be near this point
|
|
# Optimization `only_quads_related = True`:
|
|
# Avoid replicate a large self.nodeList. Get only the other elements.
|
|
# Call this in complement of `self.nodeList`
|
|
def retrieve(self, nearPoint, only_quads_related=False):
|
|
if only_quads_related:
|
|
retNodes = []
|
|
else:
|
|
retNodes = self.nodeList[:] # = list(self.nodeList)
|
|
|
|
if self.quads[0] is not None:
|
|
index = self.getIndex(nearPoint)
|
|
if index != -1:
|
|
retNodes += self.quads[index].retrieve(nearPoint)
|
|
else:
|
|
for quad in self.quads:
|
|
retNodes += quad.retrieve(nearPoint)
|
|
|
|
return retNodes
|
|
|
|
def report(self, outputArray):
|
|
if outputArray is not None:
|
|
outputArray.append(len(self.nodeList))
|
|
|
|
for quad in self.quads:
|
|
if quad is not None:
|
|
quad.report(outputArray)
|