Commit 55a83edd authored by Joseph Siddons's avatar Joseph Siddons
Browse files

fix(kdtree): increment split index if next index is above median

parent cdffc25e
...@@ -48,18 +48,32 @@ class KDTree: ...@@ -48,18 +48,32 @@ class KDTree:
points.sort(key=lambda p: getattr(p, self.variable)) points.sort(key=lambda p: getattr(p, self.variable))
split_index = n_points // 2 split_index = n_points // 2
self.partition_value = getattr(points[split_index - 1], self.variable)
while (
split_index < n_points
and getattr(points[split_index], self.variable)
== self.partition_value
):
split_index += 1
self.partition_value = getattr(points[split_index], self.variable)
self.split = True self.split = True
# Left is <= median
self.child_left = KDTree(points[:split_index], depth + 1) self.child_left = KDTree(points[:split_index], depth + 1)
# Right is > median
self.child_right = KDTree(points[split_index:], depth + 1) self.child_right = KDTree(points[split_index:], depth + 1)
return None return None
def insert(self, point: Record) -> bool: def insert(self, point: Record) -> bool:
"""Insert a Record into the KDTree. May unbalance the KDTree""" """
Insert a Record into the KDTree. May unbalance the KDTree.
The point will not be inserted if it is already in the KDTree.
"""
if not self.split: if not self.split:
if point in self.points:
return False
self.points.append(point) self.points.append(point)
return True return True
...@@ -83,7 +97,7 @@ class KDTree: ...@@ -83,7 +97,7 @@ class KDTree:
return self.child_right.delete(point) return self.child_right.delete(point)
def query(self, point) -> tuple[Record | None, float]: def query(self, point) -> tuple[Record | None, float]:
"""Find the nearest Record within the KDTree to a _query Record""" """Find the nearest Record within the KDTree to a query Record"""
if point.lon < 0: if point.lon < 0:
point2 = Record(point.lon + 360, point.lat) point2 = Record(point.lon + 360, point.lat)
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment