49 lines
1.9 KiB
Python
49 lines
1.9 KiB
Python
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
|