Add functions kNearestNeighbors and remove.

Author: ijt

Data/Trees/KdTree.hs

@@ -19,7 +19,7 @@ class Point p where
  -- |dist2 returns the squared distance between two points.
  dist2 :: p -> p -> Double
  dist2 a b = sum . map diff2 $ [0..dimension a - 1]
-where diff2 i = (coord i a - coord i b)^2
+ where diff2 i = (coord i a - coord i b)^2
 
 -- |compareDistance p a b compares the distances of a and b to p.
 compareDistance :: (Point p) => p -> p -> p -> Ordering
@@ -37,9 +37,9 @@ instance Point Point3d where
 
 
 data KdTree point = KdNode { kdLeft :: KdTree point,
- kdPoint :: point,
+  kdPoint :: point,
  kdRight :: KdTree point,
- kdAxis :: Int }
+  kdAxis :: Int }
  | KdEmpty
  deriving (Eq, Ord, Show)
 
@@ -50,8 +50,8 @@ instance Functor KdTree where
 instance F.Foldable KdTree where
  foldr f init KdEmpty = init
  foldr f init (KdNode l x r _) = F.foldr f init3 l
-where init3 = f x init2
-init2 = F.foldr f init r
+ where init3 = f x init2
+ init2 = F.foldr f init r
 
 fromList :: Point p => [p] -> KdTree p
 fromList points = fromListWithDepth points 0
@@ -62,16 +62,16 @@ fromListWithDepth [] _ = KdEmpty
 fromListWithDepth points depth = node
  where axis = axisFromDepth (head points) depth
 
- -- Sort point list and choose median as pivot element
- sortedPoints =
-L.sortBy (\a b -> coord axis a `compare` coord axis b) points
- medianIndex = length sortedPoints `div` 2
-
- -- Create node and construct subtrees
- node = KdNode { kdLeft = fromListWithDepth (take medianIndex sortedPoints) (depth+1),
- kdPoint = sortedPoints !! medianIndex,
- kdRight = fromListWithDepth (drop (medianIndex+1) sortedPoints) (depth+1),
- kdAxis = axis }
+  -- Sort point list and choose median as pivot element
+  sortedPoints =
+ L.sortBy (\a b -> coord axis a `compare` coord axis b) points
+  medianIndex = length sortedPoints `div` 2
+ 
+  -- Create node and construct subtrees
+  node = KdNode { kdLeft = fromListWithDepth (take medianIndex sortedPoints) (depth+1),
+  kdPoint = sortedPoints !! medianIndex,
+  kdRight = fromListWithDepth (drop (medianIndex+1) sortedPoints) (depth+1),
+  kdAxis = axis }
 
 axisFromDepth :: Point p => p -> Int -> Int
 axisFromDepth p depth = depth `mod` k
@@ -90,18 +90,18 @@ nearestNeighbor (KdNode KdEmpty p KdEmpty _) probe = Just p
 nearestNeighbor (KdNode l p r axis) probe =
  if xProbe <= xp then doStuff l r else doStuff r l
  where xProbe = coord axis probe
- xp = coord axis p
+  xp = coord axis p
  doStuff tree1 tree2 =
-let candidates1 = case nearestNeighbor tree1 probe of
- Nothing -> [p]
- Just best1 -> [best1, p]
- sphereIntersectsPlane = (xProbe - xp)^2 <= dist2 probe p
- candidates2 = if sphereIntersectsPlane
- then candidates1 ++ maybeToList (nearestNeighbor tree2 probe)
- else candidates1 in
-Just . L.minimumBy (compareDistance probe) $ candidates2
-
--- |invariant tells whether the KD tree property holds for a given tree and
+ let candidates1 = case nearestNeighbor tree1 probe of
+  Nothing -> [p]
+  Just best1 -> [best1, p]
+  sphereIntersectsPlane = (xProbe - xp)^2 <= dist2 probe p
+  candidates2 = if sphereIntersectsPlane
+  then candidates1 ++ maybeToList (nearestNeighbor tree2 probe)
+  else candidates1 in
+ Just . L.minimumBy (compareDistance probe) $ candidates2
+
+-- |invariant tells whether the K-D tree property holds for a given tree and
 -- all its subtrees.
 -- Specifically, it tests that all points in the left subtree lie to the left
 -- of the plane, p is on the plane, and all points in the right subtree lie to
@@ -110,16 +110,33 @@ invariant :: Point p => KdTree p -> Bool
 invariant KdEmpty = True
 invariant (KdNode l p r axis) = leftIsGood && rightIsGood
  where x = coord axis p
- leftIsGood = all ((<= x) . coord axis) (toList l)
- rightIsGood = all ((>= x) . coord axis) (toList r)
+  leftIsGood = all ((<= x) . coord axis) (toList l)
+  rightIsGood = all ((>= x) . coord axis) (toList r)
 
+-- |invariant' tells whether the K-D tree property holds for all subtrees.
 invariant' :: Point p => KdTree p -> Bool
 invariant' = all invariant . subtrees
 
+kNearestNeighbors :: (Eq p, Point p) => KdTree p -> Int -> p -> [p]
+kNearestNeighbors KdEmpty _ _ = []
+kNearestNeighbors _ k _ | k <= 0 = []
+kNearestNeighbors tree k probe = nearest : kNearestNeighbors tree' (k-1) probe
+ where nearest = fromJust $ nearestNeighbor tree probe
+ tree' = tree `remove` nearest
+
+remove :: (Eq p, Point p) => KdTree p -> p -> KdTree p
+remove KdEmpty _ = KdEmpty
+remove (KdNode l p r axis) pKill =
+ if p == pKill
+ then fromListWithDepth (toList l ++ toList r) axis
+ else if coord axis pKill <= coord axis p
+ then KdNode (remove l pKill) p r axis
+ else KdNode l p (remove r pKill) axis
+
 instance Arbitrary Point3d where
  arbitrary = do
-x <- arbitrary
-y <- arbitrary
-z <- arbitrary
-return (Point3d x y z)
+ x <- arbitrary
+ y <- arbitrary
+ z <- arbitrary
+ return (Point3d x y z)
 

KdTreeTest.hs

@@ -20,16 +20,22 @@ prop_nearestNeighbor :: [Kd.Point3d] -> Kd.Point3d -> Bool
 prop_nearestNeighbor points probe =
  Kd.nearestNeighbor tree probe == bruteNearestNeighbor points probe
  where tree = Kd.fromList points
+ bruteNearestNeighbor :: [Kd.Point3d] -> Kd.Point3d -> Maybe Kd.Point3d
+ bruteNearestNeighbor [] _ = Nothing
+ bruteNearestNeighbor points probe =
+ Just . head . L.sortBy (Kd.compareDistance probe) $ points
 
 prop_pointsAreClosestToThemselves :: [Kd.Point3d] -> Bool
 prop_pointsAreClosestToThemselves points =
  map Just points == map (Kd.nearestNeighbor tree) points
  where tree = Kd.fromList points
 
-bruteNearestNeighbor :: [Kd.Point3d] -> Kd.Point3d -> Maybe Kd.Point3d
-bruteNearestNeighbor [] _ = Nothing
-bruteNearestNeighbor points probe =
- Just . head . L.sortBy (Kd.compareDistance probe) $ points
+prop_kNearestNeighborsMatchesBrute :: [Kd.Point3d] -> Int -> Kd.Point3d -> Bool
+prop_kNearestNeighborsMatchesBrute points k p =
+ L.sort (Kd.kNearestNeighbors tree k p) == L.sort (bruteKnearestNeighbors points k p)
+ where tree = Kd.fromList points
+ bruteKnearestNeighbors points k p =
+ take k . L.sortBy (Kd.compareDistance p) $ points
 
 main = $quickCheckAll