199 lines
7.0 KiB
Haskell
199 lines
7.0 KiB
Haskell
{-# OPTIONS_HADDOCK ignore-exports #-}
|
|
|
|
module Algorithms.KDTree (kdTree
|
|
, kdFoldl
|
|
, kdFoldr
|
|
, kdTreeToRoseTree
|
|
, rangeSearch
|
|
, getValS
|
|
, isLeaf
|
|
, getVal
|
|
, getDirection
|
|
, goLeft
|
|
, goRight
|
|
, Direction(Vertical, Horizontal)
|
|
, KDTree(KTNil, KTNode))
|
|
where
|
|
|
|
|
|
import Algebra.Vector
|
|
import Data.Maybe (fromJust, catMaybes)
|
|
import Data.Tree
|
|
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 :: [P2 Double] -- ^ list of points to construct the kd-tree from
|
|
-> Direction -- ^ initial direction of the root-node
|
|
-> KDTree (P2 Double) -- ^ 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' :: P2 Double -- ^ the pivot to partition against
|
|
-> (P2 Double -> P2 Double -> Ordering) -- ^ ptCmpY or ptCmpX
|
|
-> ([P2 Double], [P2 Double]) -- ^ both lists (X, Y) or (Y, X)
|
|
-> (([P2 Double], [P2 Double]), ([P2 Double], [P2 Double])) -- ^ ((x1, x2), (y1, y2)) or
|
|
-- ((y1, y2), (x1, x2))
|
|
partition' piv cmp' (xs, ys) = ((x1, x2), (y1, y2))
|
|
where
|
|
y1 = takeWhile (/= piv) ys
|
|
y2 = tailDef [] . dropWhile (/= piv) $ ys
|
|
x1 = foldr (\x y -> [x | cmp' x piv == LT] ++ y) [] xs
|
|
x2 = foldr (\x y -> [x | cmp' x piv == GT] ++ 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 :: ([P2 Double], [P2 Double]) -- ^ both lists (X, Y)
|
|
-> (([P2 Double], [P2 Double]), ([P2 Double], [P2 Double])) -- ^ ((x1, x2), (y1, y2))
|
|
partitionY (xs, ys) = partition' (fromJust . pivot $ ys) ptCmpY (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 :: ([P2 Double], [P2 Double]) -- ^ both lists (X, Y)
|
|
-> (([P2 Double], [P2 Double]), ([P2 Double], [P2 Double])) -- ^ ((x1, x2), (y1, y2))
|
|
partitionX (xs, ys) = (\(x, y) -> (y, x))
|
|
. partition' (fromJust . pivot $ xs) ptCmpX $ (ys, xs)
|
|
|
|
|
|
-- |Execute a range search in O(log n). It returns a tuple
|
|
-- of the points found in the range and also gives back a pretty
|
|
-- rose tree suitable for printing.
|
|
rangeSearch :: KDTree (P2 Double) -- ^ tree to search in
|
|
-> ((Double, Double), (Double, Double)) -- ^ square describing the range
|
|
-> ([P2 Double], Tree String)
|
|
rangeSearch kd' sq' = (goPt kd' sq', goTree kd' sq' True)
|
|
where
|
|
-- either y1 or x1 depending on the orientation
|
|
p1' dir ((x1, y1), (_, _)) = if' (dir == Vertical) y1 x1
|
|
-- either y2 or x2 depending on the orientation
|
|
p2' dir ((_, _), (x2, y2)) = if' (dir == Vertical) y2 x2
|
|
-- either the second or first of the tuple, depending on the orientation
|
|
cur' dir = if' (dir == Vertical) snd fst
|
|
-- All points in the range.
|
|
goPt :: KDTree (P2 Double) -> ((Double, Double), (Double, Double)) -> [P2 Double]
|
|
goPt KTNil _ = []
|
|
goPt (KTNode ln pt dir rn) sq =
|
|
[pt | inRange sq pt]
|
|
++ (if' (p1' dir sq < (cur' dir . unp2 $ pt))
|
|
(goPt ln sq)
|
|
[]
|
|
)
|
|
++ (if' ((cur' dir . unp2 $ pt) < p2' dir sq)
|
|
(goPt rn sq)
|
|
[])
|
|
-- A pretty rose tree suitable for printing.
|
|
goTree :: KDTree (P2 Double) -> ((Double, Double), (Double, Double)) -> Bool -> Tree String
|
|
goTree KTNil _ _ = Node "nil" []
|
|
goTree (KTNode ln pt dir rn) sq vis
|
|
| ln == KTNil && rn == KTNil = Node treeText []
|
|
| otherwise =
|
|
Node treeText
|
|
[if' (p1' dir sq < (cur' dir . unp2 $ pt))
|
|
(goTree ln sq vis)
|
|
(goTree ln sq False)
|
|
, if' ((cur' dir . unp2 $ pt) < p2' dir sq)
|
|
(goTree rn sq vis)
|
|
(goTree rn sq False)]
|
|
where
|
|
treeText
|
|
-- HACK: in order to give specific nodes a specific color
|
|
| vis && inRange sq pt = "** " ++ (show . unp2 $ pt)
|
|
| vis = "* " ++ (show . unp2 $ pt)
|
|
| otherwise = show . unp2 $ pt
|
|
|
|
|
|
|
|
-- |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
|
|
|
|
|
|
-- |Convert a kd-tree to a rose tree, for pretty printing.
|
|
kdTreeToRoseTree :: KDTree (P2 Double) -> Tree String
|
|
kdTreeToRoseTree (KTNil) = Node "nil" []
|
|
kdTreeToRoseTree (KTNode ln val _ rn) =
|
|
Node (show . unp2 $ val) [kdTreeToRoseTree ln, kdTreeToRoseTree rn]
|
|
|
|
|
|
goLeft :: KDTree a -> Maybe (KDTree a)
|
|
goLeft (KTNode ln _ _ _) = Just ln
|
|
goLeft _ = Nothing
|
|
|
|
|
|
goRight :: KDTree a -> Maybe (KDTree a)
|
|
goRight (KTNode _ _ _ rn) = Just rn
|
|
goRight _ = Nothing
|
|
|