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 def retrieve(self, nearPoint): retNodes = 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)