module NearestNeighbour2D (KDTree, empty, fromList, findNearest) where import Data.List (sort, sortBy, minimumBy) import Test.QuickCheck import Data.Maybe (fromJust) -- | Our dimentions are called fst and snd, seeing as how we keep them in tuples data Dim = Fst | Snd deriving (Show) -- | Gets the value of a given point in the given dimention selector :: Dim -> (((a, a), b) -> a) selector Fst = fst . fst selector Snd = snd . fst -- | Updates a given point so that it's value in the given dimention equals the third argument updator :: Dim -> (a, a) -> a -> (a, a) updator Fst = (\(_, y) x -> (x, y)) updator Snd = (\(x, _) y -> (x, y)) -- | Returns a metric between two points distance :: (Num v) => (v, v) -> (v, v) -> v distance (x1, y1) (x2, y2) = a*a + b*b where a = x2 - x1 b = y2 - y1 data (Ord v, Num v) => KDTree v e = Branch Dim (v, v) e (KDTree v e) (KDTree v e) | Empty deriving (Show) -- | An empty tree. empty :: (Ord v, Num v) => KDTree v e empty = Empty -- | Construct a KD tree from a list of points. fromList :: (Ord v, Num v) => [((v, v), e)] -> KDTree v e fromList values = make dimlist values where dimlist = cycle [Fst, Snd] make _ [] = Empty make (dim:dimlist) values' = Branch dim v e (make dimlist leftvalues) (make dimlist rightvalues) where values = sortBy (\x y -> compare ((selector dim) x) ((selector dim) y)) values' midindex = length values `div` 2 (v, e) = values !! midindex (leftvalues, rightvalues') = splitAt midindex values rightvalues = if null rightvalues' then [] else tail rightvalues' -- | Find the point in the KD tree which is closest to the given point findNearest :: (Ord v, Num v, Monad m) => (v, v) -> KDTree v e -> m ((v, v), e) findNearest _ Empty = fail "Empty tree" findNearest target branch = return \$ sel \$ inner target branch where sel (x, y, _) = (x, y) inner target (Branch dim v e left right) = if radiusIntersectsFarSide then (if fardist < neardist then (farv, fare, fardist) else (nearv, neare, neardist)) else (nearv, neare, neardist) where (farv, fare, fardist) = best far (nearv, neare, neardist) = if mydist < neardist' then (v, e, mydist) else (nearv', neare', neardist') (nearv', neare', neardist') = best near mydist = distance v target best side = case side of Empty -> (v, e, mydist) otherwise -> inner target side radiusIntersectsFarSide = distance closestfarpoint target < neardist closestfarpoint = (updator dim) target \$ selector dim (v, ()) (near, far) = if (selector dim (target, ())) < (selector dim (v, ())) then (left, right) else (right, left) {- - Testing... -- | Finds the closest point in a list of points, returning that point and the -- distance from the target simpleNN :: (Ord v, Num v) => [(v, v)] -> (v, v) -> (v, v) simpleNN [] target = error "Empty list to simpleNN" simpleNN values target = fst \$ minimumBy (\x y -> compare (snd x) (snd y)) \$ map (\x -> (x, distance target x)) values prop_CorrectFloat :: [(Float, Float)] -> (Float, Float) -> Property prop_CorrectFloat xs v = not (null xs) ==> distance v a == distance v b where a = simpleNN xs v b = (fst \$ fromJust \$ findNearest v \$ fromList \$ map (\x -> (x, ())) xs) prop_CorrectInt :: [(Int, Int)] -> (Int, Int) -> Property prop_CorrectInt xs v = not (null xs) ==> distance v a == distance v b where a = simpleNN xs v b = (fst \$ fromJust \$ findNearest v \$ fromList \$ map (\x -> (x, ())) xs) runTests = do check (defaultConfig { configMaxTest = 10000, configMaxFail = 100000 }) prop_CorrectFloat check (defaultConfig { configMaxTest = 10000, configMaxFail = 100000 }) prop_CorrectInt -}