from helix.helpers.nodequadtree import Bounds, NodeQuadTree class GraphNodeStore(list): def __init__(self): super().__init__() self.variance = 0.2 self.first = True self.quadTree = None def add_node(self, node): self.append(node) if self.first: self.left = self.right = node.coordinate.x self.top = self.bottom = node.coordinate.y self.first = False else: self.left = min(self.left, node.coordinate.x) self.right = max(self.right, node.coordinate.x) self.bottom = min(self.bottom, node.coordinate.y) self.top = max(self.top, node.coordinate.y) def distance_squared(self, node, coordinate): dx = node.coordinate.x - coordinate.x dy = node.coordinate.y - coordinate.y return dx * dx + dy * dy # FIXME: This is slow if called thousands of times. Do not add any overhead in this method. def find_coordinate(self, coordinate): # create and populate the quadtree on first request if self.quadTree is None: self.quadTree = NodeQuadTree(1, Bounds(self.left, self.bottom, self.right - self.left, self.top - self.bottom), self.variance) for node in self: self.quadTree.insert(node) del self[:] variance_square = self.variance ** 2 # Optimization: avoid creating a copy of the big list `self.quadTree.nodeList` # Old: possibilities = self.quadTree.retrieve(coordinate) possibilities = self.quadTree.nodeList for node in possibilities: if self.distance_squared(node, coordinate) <= variance_square: return node more_possibilities = self.quadTree.retrieve(coordinate, only_quads_related=True) for node in more_possibilities: if self.distance_squared(node, coordinate) <= variance_square: return node return None