diff --git a/GeoSpatialTools/octtree.py b/GeoSpatialTools/octtree.py index 33e32c561a2d5951a308704457e6d658d34a7e6c..250f145b6c0061cf83b94e0901bd5ffda3f2ec17 100644 --- a/GeoSpatialTools/octtree.py +++ b/GeoSpatialTools/octtree.py @@ -461,6 +461,24 @@ class OctTree: out += f"{self.southeastfwd}" return out + def len(self, _current_len: int = 0) -> int: + """Get the number of points in the OctTree""" + _current_len += len(self.points) + print(_current_len) + if not self.divided: + return _current_len + + _current_len = self.northeastback.len(_current_len) + _current_len = self.northwestback.len(_current_len) + _current_len = self.southeastback.len(_current_len) + _current_len = self.southwestback.len(_current_len) + _current_len = self.northeastfwd.len(_current_len) + _current_len = self.northwestfwd.len(_current_len) + _current_len = self.southeastfwd.len(_current_len) + _current_len = self.southwestfwd.len(_current_len) + + return _current_len + def divide(self): """Divide the QuadTree""" self.northwestfwd = OctTree( diff --git a/GeoSpatialTools/quadtree.py b/GeoSpatialTools/quadtree.py index b7a3ad7ef470e11ea6bcda1ae0efcfa89025867b..1ef4e310a52bc75c7f8c3f7eacddfd5d595959eb 100644 --- a/GeoSpatialTools/quadtree.py +++ b/GeoSpatialTools/quadtree.py @@ -340,6 +340,19 @@ class QuadTree: out += f"{self.southeast}" return out + def len(self, _current_len: int = 0) -> int: + """Get the number of points in the OctTree""" + _current_len += len(self.points) + if not self.divided: + return _current_len + + _current_len = self.northeast.len(_current_len) + _current_len = self.northwest.len(_current_len) + _current_len = self.southeast.len(_current_len) + _current_len = self.southwest.len(_current_len) + + return _current_len + def divide(self): """Divide the QuadTree""" self.northwest = QuadTree( diff --git a/test/test_octtree.py b/test/test_octtree.py index 4a04eb25f97c8e14eb387ec2bdc1fe371df74b2a..7a3d9792294b20f9b210677ef6fd3966fce58b78 100644 --- a/test/test_octtree.py +++ b/test/test_octtree.py @@ -198,6 +198,7 @@ class TestOctTree(unittest.TestCase): ] for point in points: otree.insert(point) + assert otree.len() == len(points) - 1 # NOTE: 1 point not added assert otree.divided expected = [ points[:3], diff --git a/test/test_quadtree.py b/test/test_quadtree.py index 86a643e3d4b62b364d0e4577f59b3d2999f58aeb..9f4546c80d54e279b82686e70fd99cc4f5ae57ce 100644 --- a/test/test_quadtree.py +++ b/test/test_quadtree.py @@ -94,6 +94,7 @@ class TestQuadTree(unittest.TestCase): for point in points: qtree.insert(point) assert qtree.divided + assert qtree.len() == len(points) res = [ qtree.points, qtree.northwest.points,