Add nearestNeighbor.

Author: ijt

KdTree.hs

@@ -1,8 +1,9 @@
 {-# LANGUAGE TemplateHaskell #-}
 
--- Based on
 -- http://en.wikipedia.org/wiki/K-d_tree
--- Translated by Issac Trotts
+-- Issac Trotts
+
+import Data.Maybe
 
 import qualified Data.Foldable as F
 import qualified Data.List as L
@@ -11,11 +12,21 @@ import Test.QuickCheck
 import Test.QuickCheck.All
 
 class Point p where
+ -- |dimension returns the number of coordinates of a point.
  dimension :: p -> Int
 
- -- Get the k'th coordinate, starting from 0
+ -- |coord gets the k'th coordinate, starting from 0.
  coord :: Int -> p -> Double
 
+ -- |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
+
+-- |compareDistance p a b compares the distances of a and b to p.
+compareDistance :: (Point p) => p -> p -> p -> Ordering
+compareDistance p a b = dist2 p a `compare` dist2 p b
+
 data Point3d = Point3d { p3x :: Double, p3y :: Double, p3z :: Double }
  deriving (Eq, Ord, Show)
 
@@ -75,6 +86,23 @@ subtrees :: KdTree p -> [KdTree p]
 subtrees KdEmpty = [KdEmpty]
 subtrees t@(KdNode l x r axis) = subtrees l ++ [t] ++ subtrees r
 
+nearestNeighbor :: Point p => KdTree p -> p -> Maybe p
+nearestNeighbor KdEmpty probe = Nothing
+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
+ 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
+
 -- Testing -----
 
 -- |invariant tells whether the KD tree property holds for a given tree and
@@ -105,5 +133,15 @@ prop_invariant points = invariant' . fromList $ points
 prop_samePoints :: [Point3d] -> Bool
 prop_samePoints points = L.sort points == (L.sort . toList . fromList $ points)
 
+prop_nearestNeighbor :: [Point3d] -> Point3d -> Bool
+prop_nearestNeighbor points probe =
+ nearestNeighbor tree probe == bruteNearestNeighbor points probe
+ where tree = fromList points
+
+bruteNearestNeighbor :: [Point3d] -> Point3d -> Maybe Point3d
+bruteNearestNeighbor [] _ = Nothing
+bruteNearestNeighbor points probe =
+ Just . head . L.sortBy (compareDistance probe) $ points
+
 main = $quickCheckAll