ALGO: save direction in KDTree and refactor

This commit is contained in:
hasufell 2014-11-29 18:06:07 +01:00
parent 7dbd3a1352
commit d195d3f11d
No known key found for this signature in database
GPG Key ID: 220CD1C5BDEED020
2 changed files with 61 additions and 40 deletions

View File

@ -5,7 +5,7 @@ import Algebra.VectorTypes
import Algebra.Vector import Algebra.Vector
import Data.Maybe (fromJust, catMaybes) import Data.Maybe (fromJust, catMaybes)
import Diagrams.TwoD.Types import Diagrams.TwoD.Types
import MyPrelude (pivot) import MyPrelude (pivot,if',Not, not')
import Safe import Safe
@ -14,42 +14,35 @@ data KDTree a
-- |An empty node. -- |An empty node.
= KTNil = KTNil
-- |A node with a value and a left and right child -- |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) 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 data Direction = Vertical
| Horizontal | 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). -- |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 where
go [] _ _ = KTNil go [] _ _ = KTNil
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 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. -- |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 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). -- |Execute a range search in O(log n).
rangeSearch :: KDTree PT -> Square -> [PT] 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. -- |Left fold over ALL tree nodes.
kdFoldl :: (a -> KDTree b -> a) -> a -> KDTree b -> a 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 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. -- |Whether the tree is a leaf.
isLeaf :: KDTree a -> Bool isLeaf :: KDTree a -> Bool
isLeaf (KTNode KTNil _ KTNil) = True isLeaf (KTNode KTNil _ _ KTNil) = True
isLeaf _ = False isLeaf _ = False
-- |Get the value of the root node of the tree. Returns Nothing if it's a -- |Get the value of the root node of the tree. Returns Nothing if it's a
-- leaf. -- leaf.
getVal :: KDTree a -> Maybe a getVal :: KDTree a -> Maybe a
getVal (KTNode _ val _) = Just val getVal (KTNode _ val _ _) = Just val
getVal _ = Nothing getVal _ = Nothing
-- |Get the direction of the current node/level.
getDirection :: KDTree a -> Maybe Direction
getDirection (KTNode _ _ dir _) = Just dir
getDirection _ = Nothing

View File

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