Fixing behavior of KDTree.remove() when an element occurs more than o…

…nce within
LuizZak · Jul 24, 2024 · 3577f52 · 3577f52
1 parent 46863ad
commit 3577f52
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/Sources/GeometriaAlgorithms/SpatialPartitioning/KDTree/KDTree.swift b/Sources/GeometriaAlgorithms/SpatialPartitioning/KDTree/KDTree.swift
@@ -137,19 +137,15 @@ public struct KDTree<Element: KDTreeLocatable> where Element.Vector: VectorCompa
     public mutating func remove(_ element: Element) where Element: Equatable {
         ensureUnique()
 
-        guard let root else {
-            return
-        }
-
         var allNearest: [Subdivision] = []
-        root.nearestSubdivisions(
+        root?.nearestSubdivisions(
             to: element.location,
             distanceSquared: .zero,
             onMatch: { if $0.element == element { allNearest.append($0) } }
         )
 
         for nearest in allNearest {
-            guard let path = root.findPath(to: nearest) else {
+            guard let path = self.root?.findPath(to: nearest) else {
                 return
             }
 

diff --git a/Tests/GeometriaAlgorithmsTests/SpatialPartitioning/KDTree/KDTreeTests.swift b/Tests/GeometriaAlgorithmsTests/SpatialPartitioning/KDTree/KDTreeTests.swift
@@ -113,6 +113,23 @@ class KDTreeTests: XCTestCase {
         }
     }
 
+    func testRemove_repeatedElement() {
+        let points = makePointCloud(count: 50)
+        var sut = KDTree(elements: points)
+        for point in points {
+            sut.insert(point)
+        }
+
+        for point in points[..<25] {
+            sut.remove(point)
+        }
+
+        assertIsValid(sut)
+        for point in sut.elements() {
+            XCTAssertEqual(sut.nearestNeighbor(to: point), point)
+        }
+    }
+
     func testNearestNeighbor_pointCloudPartition_2D_identitySearch() {
         let points = makePointCloud(count: 50)
         let sut = KDTree<Vector2D>(elements: points)