ALGO: initial implementation of KDTree

No visualization yet.

Algorithms/KDTree/KDTree.hs

@@ -0,0 +1,116 @@
+module Algorithms.KDTree.KDTree where
+
+
+import Algebra.VectorTypes
+import Algebra.Vector
+import Data.Maybe (fromJust, catMaybes)
+import Diagrams.TwoD.Types
+import MyPrelude (pivot)
+import Safe
+
+
+-- |The KDTree data structure.
+data KDTree a
+  -- |An empty node.
+  = KTNil
+  -- |A node with a value and a left and right child
+  | KTNode (KDTree a) a (KDTree a)
+  deriving (Show, Eq)
+
+data Crumb a = Left (KDTree a)
+             | Right (KDTree a)
+  deriving (Show, Eq)
+
+-- |A list of Crumbs.
+type Breadcrumbs a = [Crumb a]
+
+-- |Zipper for the KDTree.
+type Zipper a = (KDTree a, Breadcrumbs a)
+
+data Direction = Vertical
+               | Horizontal
+
+
+-- |Construct a kd-tree from a list of points in O(n log n).
+kdTree :: [PT]
+       -> KDTree PT
+kdTree xs' = go (sortedX xs') (sortedY xs') Horizontal
+  where
+    go [] _ _ = KTNil
+    go _ [] _ = KTNil
+    go xs ys Vertical =
+      KTNode (go x1 y1 Horizontal)
+             (fromJust . pivot $ ys)
+             (go x2 y2 Horizontal)
+      where
+        ((x1, x2), (y1, y2)) = partition' (fromJust . pivot $ ys) (xs, ys)
+    go xs ys Horizontal =
+      KTNode (go x1 y1 Vertical)
+             (fromJust . pivot $ xs)
+             (go x2 y2 Vertical)
+      where
+        ((y1, y2), (x1, x2)) = partition' (fromJust . pivot $ xs) (ys, xs)
+
+
+-- |Partitions two sorted list of points X and Y against a pivot.
+-- If you want to partition against the pivot of Y, then you pass
+--   partition' (pivot ys) (xs, ys)
+-- and get ((x1, x2), (y1, y2)).
+-- If you want to partition against the pivot of X, then you pass
+--   partition' (pivot xs) (ys, xs)
+-- and get ((y1, y2), (x1, x2)).
+partition' :: PT                           -- ^ the pivot to partition against
+           -> ([PT], [PT])                 -- ^ both lists (X, Y) or (Y, X)
+           -> (([PT], [PT]), ([PT], [PT])) -- ^ ((x1, x2), (y1, y2)) or
+                                           --   ((y1, y2), (x1, x2))
+partition' piv (xs, ys) = ((x1, x2), (y1, y2))
+  where
+    y1 = takeWhile (/= piv) ys
+    y2 = tailDef [] . dropWhile (/= piv) $ ys
+    x1 = foldr (\x y -> [x | x `elem` y1] ++ y) [] xs
+    x2 = foldr (\x y -> [x | x `elem` y2] ++ y) [] xs
+
+
+-- |Execute a range search in O(log n).
+rangeSearch :: KDTree PT -> Square -> [PT]
+rangeSearch = go Horizontal
+  where
+    go _ KTNil _ = []
+    go Vertical (KTNode ln pt rn) sq@(_, (y1, y2)) =
+      [pt | inRange sq pt]
+      ++ (if y1 < (snd . unp2 $ pt) then go Horizontal ln sq else [])
+      ++ (if (snd . unp2 $ pt) < y2 then go Horizontal rn sq else [])
+    go Horizontal (KTNode ln pt rn) sq@((x1, x2), _) =
+      [pt | inRange sq pt]
+      ++ (if x1 < (fst . unp2 $ pt) then go Vertical ln sq else [])
+      ++ (if (fst . unp2 $ pt) < x2 then go Vertical rn sq else [])
+
+
+-- |Left fold over ALL tree nodes.
+kdFoldl :: (a -> KDTree b -> a) -> a -> KDTree b -> a
+kdFoldl f sv kd@(KTNode ln _ rn) = foldl (kdFoldl f) (f sv kd) [ln, rn]
+kdFoldl f sv kd = f sv kd
+
+
+-- |Right fold over ALL tree nodes.
+kdFoldr :: (KDTree b -> a -> a) -> a -> KDTree b -> a
+kdFoldr f sv kd = kdFoldl (\g b x -> g (f b x)) id kd sv
+
+
+-- |Get all values of a tree.
+getValS :: KDTree a -> [a]
+getValS = catMaybes . kdFoldl (\x y -> x ++ [getVal y]) []
+
+
+-- |Whether the tree is a leaf.
+isLeaf :: KDTree a -> Bool
+isLeaf (KTNode KTNil _ KTNil) = True
+isLeaf _                      = False
+
+
+-- |Get the value of the root node of the tree. Returns Nothing if it's a
+-- leaf.
+getVal :: KDTree a -> Maybe a
+getVal (KTNode _ val _) = Just val
+getVal _                = Nothing
+

CG2.cabal

@@ -58,6 +58,7 @@ executable Gtk
                        Algorithms.ConvexHull.GrahamScan
                        Algorithms.PolygonIntersection.Core
                        Algorithms.QuadTree.QuadTree
+                       Algorithms.KDTree.KDTree
                        Graphics.Diagram.Gtk
                        Graphics.Diagram.Plotter
                        Graphics.Diagram.Types
@@ -104,6 +105,7 @@ executable Gif
                        Algorithms.ConvexHull.GrahamScan
                        Algorithms.PolygonIntersection.Core
                        Algorithms.QuadTree.QuadTree
+                       Algorithms.KDTree.KDTree
                        Graphics.Diagram.Gif
                        Graphics.Diagram.Plotter
                        Graphics.Diagram.Types