kd-tree (empty) → 0.1.0

raw patch · 4 files changed

+173/−0 lines, 4 filesdep +basedep +lensdep +linearsetup-changed

Dependencies added: base, lens, linear, vector, vector-algorithms

Files

LICENSE +30/−0
Setup.hs +2/−0
kd-tree.cabal +28/−0
src/Data/KdTree.hs +113/−0

+ LICENSE view

@@ -0,0 +1,30 @@+Copyright (c) 2014, Ben Gamari++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++    * Redistributions of source code must retain the above copyright+      notice, this list of conditions and the following disclaimer.++    * Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++    * Neither the name of Ben Gamari nor the names of other+      contributors may be used to endorse or promote products derived+      from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

+ Setup.hs view

@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain

+ kd-tree.cabal view

@@ -0,0 +1,28 @@+name:                kd-tree+version:             0.1.0+synopsis:            A simple k-d tree implementation+description:         A simple k-d tree implementation+homepage:            http://github.com/bgamari/kd-tree+license:             BSD3+license-file:        LICENSE+author:              Ben Gamari+maintainer:          bgamari@gmail.com+copyright:           (c) 2014 Ben Gamari <bgamari@gmail.com>+category:            Data+build-type:          Simple+cabal-version:       >=1.10++source-repository head+  type:                git+  location:            git@github.com:bgamari/kd-tree++library+  exposed-modules:     Data.KdTree+  other-extensions:    ScopedTypeVariables+  build-depends:       base >=4.7 && <4.8,+                       linear >=1.10 && <1.11,+                       lens >=4.2 && <4.3,+                       vector >=0.10 && <0.11,+                       vector-algorithms >=0.6 && <0.7+  hs-source-dirs:      src+  default-language:    Haskell2010

+ src/Data/KdTree.hs view

@@ -0,0 +1,113 @@+{-# LANGUAGE ScopedTypeVariables #-}++module Data.KdTree+    ( KdTree+      -- * Construction+    , fromVector+      -- * Queries+    , nearest+    , toList+      -- * Diagnostics+    , isValid+    , showKdTree+    ) where++import Prelude hiding (sort)+import Data.List (minimumBy)+import Data.Maybe (maybeToList)+import Data.Ord (comparing)++import Linear hiding (point)+import Control.Lens+import qualified Data.Vector.Generic as V+import Data.Vector.Algorithms.Intro (sortBy)++-- | The k-d tree is a data structure capable of efficiently answering+-- nearest neighbor search queries in low-dimensional spaces. As a rule+-- of thumb, for efficient lookups the number of points in @k@ dimensions+-- should greatly exceed @2^k@+data KdTree f a = KdNode { point :: !(f a)+                         , axis  :: E f+                         , left  :: KdTree f a+                         , right :: KdTree f a+                         }+                | KdEmpty++-- | Construct a @KdTree@ from a vector of points+fromVector :: (Ord a, V.Vector v (f a)) => [E f] -> v (f a) -> KdTree f a+fromVector basis pts = go (cycle basis) pts+  where+    go _ pts | V.null pts = KdEmpty+    go (axis:rest) pts =+      let pts' = V.modify (sortBy $ comparing (^. el axis)) pts+          pivotIdx = V.length pts' `div` 2+      in KdNode { point = pts' V.! pivotIdx+                , axis  = axis+                , left  = go rest $ V.take pivotIdx pts'+                , right = go rest $ V.drop (pivotIdx+1) pts'+                }++quadranceTo :: (Num a, Metric f) => f a -> f a -> a+quadranceTo a b = quadrance (a ^-^ b)++-- | Find the point nearest to the given point+nearest :: forall f a. (Ord a, Num a, Metric f)+        => f a -> KdTree f a -> Maybe (f a)+nearest pt tree = go tree+  where+    go :: KdTree f a -> Maybe (f a)+    go KdEmpty = Nothing+    go (KdNode nodePt axis l r)+      | (pt ^. el axis) <= (nodePt ^. el axis) = go' nodePt axis l r+      | otherwise                              = go' nodePt axis r l++    go' :: f a   -- ^ The point of the node we are sitting at+        -> E f   -- ^ The splitting axis of the node+        -> KdTree f a -- ^ The subnode the query point sits in+        -> KdTree f a -- ^ The other subnode+        -> Maybe (f a)+    go' nodePt axis side other =+      let best = case go side of+                   Nothing    -> [nodePt]+                   Just best' -> [best', nodePt]+          tryAdj = (pt^.el axis - nodePt^.el axis)^2 <= quadrance (pt ^-^ nodePt)+          bestAdj = if tryAdj+                      then maybeToList $ go other+                      else []+      in Just $ minimumBy (comparing $ quadranceTo pt) (best ++ bestAdj)++-- | List all points in a tree+toList :: KdTree f a -> [f a]+toList KdEmpty = []+toList (KdNode point _ l r) = point : (toList l ++ toList r)++-- | Verify that the node is well-formed+nodeIsValid :: Ord a => KdTree f a -> Bool+nodeIsValid KdEmpty = True+nodeIsValid (KdNode point axis l r) =+       all (\p->p^.el axis <= point^.el axis) (toList l)+    && all (\p->p^.el axis >  point^.el axis) (toList r)++-- | Verify that the tree is well-formed (recursively)+isValid :: Ord a => KdTree f a -> Bool+isValid KdEmpty = True+isValid node@(KdNode _ _ l r) =+    nodeIsValid node && isValid l && isValid r++onAxis :: E f -> (a -> a -> b) -> f a -> f a -> b+onAxis (E l) f a b = f (a ^. l) (b ^. l)++-- | Given names for the axes show the tree+showKdTree :: Show (f a) => f String -> KdTree f a -> String+showKdTree axisNames tree = unlines $ fmt 0 tree+  where+    --fmt :: Int -> Kdtree f a -> [String]+    fmt depth node =+      case node of+        KdEmpty -> [indent "KdEmpty"]+        (KdNode point axis l r) ->+          [ indent $ "KdNode ("++show point++") "++show (axisNames ^. el axis) ]+          ++ fmt (depth+2) l+          ++ [""]+          ++ fmt (depth+2) r+      where indent = (replicate depth ' ' ++)