{-# LANGUAGE RecordWildCards #-}
module Data.BalBST where
import Control.Applicative((<|>))
import Data.Bifunctor
import Data.Function (on)
import Data.Functor.Contravariant
import qualified Data.List as L
import Data.Maybe
import qualified Data.Tree as T
import Prelude hiding (lookup,null)
--------------------------------------------------------------------------------
-- | Describes how to search in a tree
data TreeNavigator k a = Nav { goLeft :: a -> k -> Bool
, extractKey :: a -> a -> k
}
instance Contravariant (TreeNavigator k) where
contramap f (Nav gL eK) = Nav (\a k -> gL (f a) k) (\x y -> eK (f x) (f y))
ordNav :: Ord a => TreeNavigator a a
ordNav = Nav (<=) min
ordNavBy :: Ord b => (a -> b) -> TreeNavigator b a
ordNavBy f = Nav (\x k -> f x <= k) (min `on` f)
-- instance Functor (TreeNavigator k) where
-- fmap f Nav{..} = Nav (\b k -> )
-- | A balanced binary search tree
data BalBST k a = BalBST { nav :: !(TreeNavigator k a)
, toTree :: !(Tree k a)
}
instance (Show k, Show a) => Show (BalBST k a) where
show (BalBST _ t) = "BalBST (" ++ show t ++ ")"
data Color = Red | Black deriving (Show,Read,Eq,Ord)
type Height = Int
-- Red-Black tree with values in the leaves
data Tree k a = Empty
| Leaf !a
| Node !Color !Height (Tree k a) !k (Tree k a) deriving (Show,Eq)
--------------------------------------------------------------------------------
-- | Creates an empty BST
empty :: TreeNavigator k a -> BalBST k a
empty n = BalBST n Empty
-- | \(O(n\log n)\)
fromList :: TreeNavigator k a -> [a] -> BalBST k a
fromList n = foldr insert (empty n)
fromList' :: Ord a => [a] -> BalBST a a
fromList' = fromList ordNav
-- -- | \(O(n)\)
-- fromAscList :: TreeNavigator k a -> [a] -> BalBST k a
-- fromAscList = undefined
--------------------------------------------------------------------------------
-- | Check if the tree is empty
null :: BalBST k a -> Bool
null (BalBST _ Empty) = True
null _ = False
-- | Test if an element occurs in the BST.
-- \(O(\log n)\)
lookup :: Eq a => a -> BalBST k a -> Maybe a
lookup x (BalBST Nav{..} t) = lookup' t
where
lookup' Empty = Nothing
lookup' (Leaf y) = if x == y then Just y else Nothing
lookup' (Node _ _ l k r)
| goLeft x k = lookup' l
| otherwise = lookup' r
-- | \(O(\log n)\)
member :: Eq a => a -> BalBST k a -> Bool
member x = isJust . lookup x
-- | Search for the Predecessor
-- \(O(\log n)\)
lookupLE :: Ord k => k -> BalBST k a -> Maybe a
lookupLE kx (BalBST n@Nav{..} t) = lookup' t
where
lookup' Empty = Nothing
lookup' (Leaf y) = if goLeft y kx then Just y else Nothing
lookup' (Node _ _ l k r)
| kx <= k = lookup' l
| otherwise = lookup' r <|> lookupMax (BalBST n l)
-- | Insert an element in the BST.
--
-- \(O(\log n)\)
insert :: a -> BalBST k a -> BalBST k a
insert x (BalBST n@Nav{..} t) = BalBST n (blacken $ insert' t)
where
insert' Empty = Leaf x
insert' (Leaf y) = let k = extractKey x y
(l,r) = if goLeft x k then (x,y) else (y,x)
in red 2 (Leaf l) k (Leaf r)
insert' (Node c h l k r)
| goLeft x k = balance c h (insert' l) k r
| otherwise = balance c h l k (insert' r)
-- delete = undefined
-- | Delete (one occurance of) an element.
-- \(O(\log n)\)
delete :: Eq a => a -> BalBST k a -> BalBST k a
delete x t = let Split l _ r = split x t
n = nav t
in BalBST n $ joinWith n l r
-- (BalBST n@Nav{..} t) = delete' t
-- where
-- delete' Empty = Empty
-- delete' l@(Leaf y) = if x == y then Empty else l
-- delete' (Node c h l k r)
-- | goLeft x k =
--------------------------------------------------------------------------------
-- | Extract the minimum from the tree
-- \(O(\log n)\)
minView :: BalBST k a -> Maybe (a, Tree k a)
minView (BalBST n t) = minView' t
where
minView' Empty = Nothing
minView' (Leaf x) = Just (x,Empty)
minView' (Node _ _ l _ r) = fmap (flip (joinWith n) r) <$> minView' l
lookupMin :: BalBST k b -> Maybe b
lookupMin = fmap fst . maxView
-- | Extract the maximum from the tree
-- \(O(\log n)\)
maxView :: BalBST k a -> Maybe (a, Tree k a)
maxView (BalBST n t) = maxView' t
where
maxView' Empty = Nothing
maxView' (Leaf x) = Just (x,Empty)
maxView' (Node _ _ l _ r) = fmap (joinWith n l) <$> maxView' r
lookupMax :: BalBST k b -> Maybe b
lookupMax = fmap fst . maxView
-- | Joins two BSTs. Assumes that the ranges are disjoint. It takes the left Tree nav
--
-- \(O(\log n)\)
join :: BalBST k a -> BalBST k a -> BalBST k a
join (BalBST n l) (BalBST _ r) = BalBST n $ joinWith n l r
-- | Joins two BSTs' with a specific Tree Navigator
--
-- \(O(\log n)\)
joinWith :: TreeNavigator k a -> Tree k a -> Tree k a -> Tree k a
joinWith Nav{..} tl tr
| lh >= rh = blacken $ joinL tl tr
| otherwise = blacken $ joinR tl tr
where
rh = height tr
lh = height tl
joinL Empty _ = Empty
joinL l Empty = l
joinL l@(Leaf x) r@(Leaf y) = red 2 l (extractKey x y) r
joinL l@(Node c h ll k lr) r
| h == rh = let lm = unsafeMax lr
rm = unsafeMin r
in balance Red (h+1) l (extractKey lm rm) r
| otherwise = balance c h ll k (joinL lr r)
-- lh >= rh
joinL _ _ = error "joinL. absurd"
joinR _ Empty = Empty
joinR Empty r = r
joinR l@(Leaf x) r@(Leaf y) = red 2 l (extractKey x y) r
joinR l r@(Node c h rl k rr)
| h == lh = let lm = unsafeMax l
rm = unsafeMin rl
in balance Red (h+1) l (extractKey lm rm) r
| otherwise = balance c h (joinR l rl) k rr
-- lh >= rh
joinR _ _ = error "joinR absurd"
--------------------------------------------------------------------------------
-- | Splitting and extracting
-- | A pair that is strict in its first argument and lazy in the second.
data Pair a b = Pair { fst' :: !a
, snd' :: b
} deriving (Show,Eq,Functor,Foldable,Traversable)
collect :: b -> [Pair a b] -> Pair [a] b
collect def [] = Pair [] def
collect _ xs = Pair (map fst' xs) (snd' $ last xs)
-- | Extract a prefix from the tree, i.e. a repeated 'minView'
--
-- \(O(\log n +k)\), where \(k\) is the size of the extracted part
extractPrefix :: BalBST k a -> [Pair a (Tree k a)]
extractPrefix (BalBST n@Nav{..} t) = extractPrefix' t
where
extractPrefix' Empty = []
extractPrefix' (Leaf x) = [Pair x Empty]
extractPrefix' (Node _ _ l _ r) = ls ++ extractPrefix' r
where
ls = map (fmap $ flip (joinWith n) r) $ extractPrefix' l
-- | Extract a suffix from the tree, i.e. a repeated 'minView'
--
-- \(O(\log n +k)\), where \(k\) is the size of the extracted part
extractSuffix :: BalBST k a -> [Pair a (Tree k a)]
extractSuffix (BalBST n@Nav{..} t) = extract t
where
extract Empty = []
extract (Leaf x) = [Pair x Empty]
extract (Node _ _ l _ r) = rs ++ extract l
where
rs = map (fmap $ joinWith n l) $ extract r
-- | Result of splititng a tree
data Split a b = Split a !b a deriving (Show,Eq)
-- | Splits the tree at x. Note that if x occurs more often, no guarantees are
-- given which one is found.
--
-- \(O(\log n)\)
split :: Eq a => a -> BalBST k a -> Split (Tree k a) (Maybe a)
split x (BalBST n@Nav{..} t) = split' t
where
split' Empty = Split Empty Nothing Empty
split' l@(Leaf y)
| x == y = Split Empty (Just y) Empty
| goLeft x (extractKey x y) = Split l Nothing Empty
| otherwise = Split Empty Nothing l
split' (Node _ _ l k r)
| goLeft x k = let Split l' mx r' = split' l
in Split l' mx (joinWith n r' r)
| otherwise = let Split l' mx r' = split' r
in Split (joinWith n l l') mx r'
-- | split based on a monotonic predicate
--
-- \(O(\log n)\)
splitMonotone :: (a -> Bool) -> BalBST k a
-> (BalBST k a, BalBST k a)
splitMonotone p (BalBST n@Nav{..} t) = bimap (BalBST n) (BalBST n) $ split' t
where
split' Empty = (Empty,Empty)
split' l@(Leaf y)
| p y = (Empty,l)
| otherwise = (l,Empty)
split' (Node _ _ l _ r)
| p (unsafeMin r) = let (l',m) = split' l in (l',joinWith n m r)
| otherwise = let (m,r') = split' r in (joinWith n l m, r')
-- | Splits at a given monotone predicate p, and then selects everything that
-- satisfies the predicate sel.
splitExtract :: (a -> Bool) -> (a -> Bool) -> BalBST k a
-> Split (BalBST k a) ([a],[a])
splitExtract p sel bst = Split (BalBST n before) (reverse mid1,mid2) (BalBST n after)
where
n = nav bst
(before',after') = splitMonotone p bst
extract def = collect def . L.takeWhile (sel . fst')
Pair mid1 before = extract (toTree before') $ extractSuffix before'
Pair mid2 after = extract (toTree after') $ extractPrefix after'
--------------------------------------------------------------------------------
data T k a = Internal !Color !Height !k | Val !a deriving (Show,Eq,Ord)
toRoseTree :: Tree k a -> Maybe (T.Tree (T k a))
toRoseTree Empty = Nothing
toRoseTree (Leaf x) = Just $ T.Node (Val x) []
toRoseTree (Node c h l k r) = Just $ T.Node (Internal c h k) (mapMaybe toRoseTree [l,r])
showTree :: (Show k, Show a) => BalBST k a -> String
showTree = maybe "Empty" T.drawTree . fmap (fmap show) . toRoseTree . toTree
-- | Get the minimum in the tree. Errors when the tree is empty
--
-- \(O(\log n)\)
unsafeMin :: Tree k a -> a
unsafeMin (Leaf x) = x
unsafeMin (Node _ _ l _ _) = unsafeMin l
unsafeMin _ = error "unsafeMin: Empty"
-- | Get the maximum in the tree. Errors when the tree is empty
--
-- \(O(\log n)\)
unsafeMax :: Tree k a -> a
unsafeMax (Leaf x) = x
unsafeMax (Node _ _ _ _ r) = unsafeMax r
unsafeMax _ = error "unsafeMax: Empty"
-- | Extract all elements in the tree
--
-- \(O(n)\)
toList :: BalBST k a -> [a]
toList = toList' . toTree
-- | Extract all elements in the tree
--
-- \(O(n)\)
toList' :: Tree k a -> [a]
toList' Empty = []
toList' (Leaf x) = [x]
toList' (Node _ _ l _ r) = toList' l ++ toList' r
--------------------------------------------------------------------------------
-- * Helper stuff
black :: Height -> Tree k a -> k -> Tree k a -> Tree k a
black = Node Black
red :: Height -> Tree k a -> k -> Tree k a -> Tree k a
red = Node Red
blacken :: Tree k a -> Tree k a
blacken (Node Red h l k r) = Node Black h l k r
blacken t = t
-- | rebalance the tree
balance :: Color -> Height -> Tree k a -> k -> Tree k a -> Tree k a
balance Black h (Node Red _ (Node Red _ a x b) y c) z d = mkNode h a x b y c z d
balance Black h (Node Red _ a x (Node Red _ b y c)) z d = mkNode h a x b y c z d
balance Black h a x (Node Red _ (Node Red _ b y c) z d) = mkNode h a x b y c z d
balance Black h a x (Node Red _ b y (Node Red _ c z d)) = mkNode h a x b y c z d
balance co h a x b = Node co h a x b
mkNode :: Height
-> Tree k a -> k -> Tree k a -> k -> Tree k a -> k -> Tree k a
-> Tree k a
mkNode h a x b y c z d = red h (black h a x b) y (black h c z d)
height :: Tree k a -> Height
height Empty = 0
height (Leaf _) = 1
height (Node _ h _ _ _) = h