ALGO: save direction in KDTree and refactor

Algorithms/KDTree/KDTree.hs

@@ -5,7 +5,7 @@ import Algebra.VectorTypes
 import Algebra.Vector
 import Data.Maybe (fromJust, catMaybes)
 import Diagrams.TwoD.Types
-import MyPrelude (pivot)
+import MyPrelude (pivot,if',Not, not')
 import Safe
 
 
@@ -14,42 +14,35 @@ data KDTree a
   -- |An empty node.
   = KTNil
   -- |A node with a value and a left and right child
-  | KTNode (KDTree a) a (KDTree a)
+  | KTNode (KDTree a) a Direction (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
+  deriving (Show, Eq, Enum)
+
+instance Not Direction where
+  not' Vertical   = Horizontal
+  not' Horizontal = Vertical
 
 
 -- |Construct a kd-tree from a list of points in O(n log n).
-kdTree :: [PT]
+kdTree :: [PT]      -- ^ list of points to construct the kd-tree from
-       -> KDTree PT
+       -> Direction -- ^ initial direction of the root-node
-kdTree xs' = go (sortedX xs') (sortedY xs') Horizontal
+       -> KDTree PT -- ^ resulting kd-tree
+kdTree xs' = go (sortedX xs') (sortedY xs')
   where
     go [] _ _ = KTNil
     go _ [] _ = KTNil
-    go xs ys Vertical =
+    go xs ys dir =
-      KTNode (go x1 y1 Horizontal)
+      KTNode (go x1 y1 (not' dir))
-             (fromJust . pivot $ ys)
+             (fromJust . pivot $ if' (dir == Vertical) ys xs)
-             (go x2 y2 Horizontal)
+             dir
+             (go x2 y2 (not' dir))
       where
-        ((x1, x2), (y1, y2)) = partition' (fromJust . pivot $ ys) (xs, ys)
+        ((x1, x2), (y1, y2)) = if' (dir == Vertical)
-    go xs ys Horizontal =
+                                   (partitionY (xs, ys))
-      KTNode (go x1 y1 Vertical)
+                                   (partitionX (xs, ys))
-             (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.
@@ -71,24 +64,40 @@ partition' piv (xs, ys) = ((x1, x2), (y1, y2))
     x2 = foldr (\x y -> [x | x `elem` y2] ++ y) [] xs
 
 
+-- |Partition two sorted lists of points X and Y against the pivot of
+-- Y. This function is unsafe as it does not check if there is a valid
+-- pivot.
+partitionY :: ([PT], [PT])                 -- ^ both lists (X, Y)
+           -> (([PT], [PT]), ([PT], [PT])) -- ^ ((x1, x2), (y1, y2))
+partitionY (xs, ys) = partition' (fromJust . pivot $ ys) (xs, ys)
+
+
+-- |Partition two sorted lists of points X and Y against the pivot of
+-- X. This function is unsafe as it does not check if there is a valid
+-- pivot.
+partitionX :: ([PT], [PT])                 -- ^ both lists (X, Y)
+           -> (([PT], [PT]), ([PT], [PT])) -- ^ ((x1, x2), (y1, y2))
+partitionX (xs, ys) = (\(x, y) -> (y, x))
+                      . partition' (fromJust . pivot $ xs) $ (ys, xs)
+
+
 -- |Execute a range search in O(log n).
 rangeSearch :: KDTree PT -> Square -> [PT]
-rangeSearch = go Horizontal
+rangeSearch KTNil _ = []
-  where
+rangeSearch (KTNode ln pt Vertical rn) sq@(_, (y1, y2)) =
-    go _ KTNil _ = []
+    [pt | inRange sq pt]
-    go Vertical (KTNode ln pt rn) sq@(_, (y1, y2)) =
+    ++ (if y1 < (snd . unp2 $ pt) then rangeSearch ln sq else [])
-      [pt | inRange sq pt]
+    ++ (if (snd . unp2 $ pt) < y2 then rangeSearch rn sq else [])
-      ++ (if y1 < (snd . unp2 $ pt) then go Horizontal ln sq else [])
+rangeSearch (KTNode ln pt Horizontal rn) sq@((x1, x2), _) =
-      ++ (if (snd . unp2 $ pt) < y2 then go Horizontal rn sq else [])
+    [pt | inRange sq pt]
-    go Horizontal (KTNode ln pt rn) sq@((x1, x2), _) =
+    ++ (if x1 < (fst . unp2 $ pt) then rangeSearch ln sq else [])
-      [pt | inRange sq pt]
+    ++ (if (fst . unp2 $ pt) < x2 then rangeSearch rn sq else [])
-      ++ (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@(KTNode ln _ _ rn) = foldl (kdFoldl f) (f sv kd) [ln, rn]
 kdFoldl f sv kd = f sv kd
 
 
@@ -104,13 +113,20 @@ getValS = catMaybes . kdFoldl (\x y -> x ++ [getVal y]) []
 
 -- |Whether the tree is a leaf.
 isLeaf :: KDTree a -> Bool
-isLeaf (KTNode KTNil _ KTNil) = True
+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 (KTNode _ val _ _) = Just val
 getVal _                = Nothing
 
+
+-- |Get the direction of the current node/level.
+getDirection :: KDTree a -> Maybe Direction
+getDirection (KTNode _ _ dir _) = Just dir
+getDirection _                  = Nothing
+
+

MyPrelude.hs

@@ -10,6 +10,11 @@ class Def a where
   def :: a
 
 
+-- |For negating random types.
+class Not b where
+  not' :: b -> b
+
+
 -- |Split an array into subarrays depending on a given condition.
 splitBy :: (a -> Bool) -- ^ condition
         -> [a]         -- ^ array to split