133 lines
4.3 KiB
Haskell
133 lines
4.3 KiB
Haskell
module Algorithms.KDTree.KDTree where
|
|
|
|
|
|
import Algebra.VectorTypes
|
|
import Algebra.Vector
|
|
import Data.Maybe (fromJust, catMaybes)
|
|
import Diagrams.TwoD.Types
|
|
import MyPrelude (pivot,if',Not, not')
|
|
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 Direction (KDTree a)
|
|
deriving (Show, Eq)
|
|
|
|
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] -- ^ list of points to construct the kd-tree from
|
|
-> Direction -- ^ initial direction of the root-node
|
|
-> KDTree PT -- ^ resulting kd-tree
|
|
kdTree xs' = go (sortedX xs') (sortedY xs')
|
|
where
|
|
go [] _ _ = KTNil
|
|
go _ [] _ = KTNil
|
|
go xs ys dir =
|
|
KTNode (go x1 y1 (not' dir))
|
|
(fromJust . pivot $ if' (dir == Vertical) ys xs)
|
|
dir
|
|
(go x2 y2 (not' dir))
|
|
where
|
|
((x1, x2), (y1, y2)) = if' (dir == Vertical)
|
|
(partitionY (xs, ys))
|
|
(partitionX (xs, ys))
|
|
|
|
|
|
-- |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
|
|
|
|
|
|
-- |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 KTNil _ = []
|
|
rangeSearch (KTNode ln pt Vertical rn) sq@(_, (y1, y2)) =
|
|
[pt | inRange sq pt]
|
|
++ (if y1 < (snd . unp2 $ pt) then rangeSearch ln sq else [])
|
|
++ (if (snd . unp2 $ pt) < y2 then rangeSearch rn sq else [])
|
|
rangeSearch (KTNode ln pt Horizontal rn) sq@((x1, x2), _) =
|
|
[pt | inRange sq pt]
|
|
++ (if x1 < (fst . unp2 $ pt) then rangeSearch ln sq else [])
|
|
++ (if (fst . unp2 $ pt) < x2 then rangeSearch 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
|
|
|
|
|
|
-- |Get the direction of the current node/level.
|
|
getDirection :: KDTree a -> Maybe Direction
|
|
getDirection (KTNode _ _ dir _) = Just dir
|
|
getDirection _ = Nothing
|
|
|
|
|