module Game.Antisplice.Utils.AVL (avlMax,avlMin,avlLookup,avlHeight,avlSize,avlInsert,avlRemove,AVL (EmptyAVL,AVL),avlRoot,avlPreorder,avlPostorder,avlInorder) where
import Game.Antisplice.Utils.BST
instance Indexable i o v => AnyBST AVL i o v where
anyBstMax = avlMax
anyBstMin = avlMin
anyBstLookup = avlLookup
anyBstEmpty = EmptyAVL
anyBstInsert = avlInsert
anyBstRemove = avlRemove
data AVL a = EmptyAVL | AVL a Int Int !(AVL a) !(AVL a)
avlMax :: AVL i -> Maybe i
avlMax EmptyAVL = Nothing
avlMax (AVL a _ _ _ EmptyAVL) = Just a
avlMax (AVL _ _ _ _ r) = avlMax r
avlMin :: AVL i -> Maybe i
avlMin EmptyAVL = Nothing
avlMin (AVL a _ _ EmptyAVL _) = Just a
avlMin (AVL _ _ _ l _) = avlMin l
avlLookup :: Indexable i o v => o -> AVL i -> Maybe v
avlLookup _ EmptyAVL = Nothing
avlLookup o (AVL a _ _ l r)
| o == indexOf a = Just $ valueOf a
| o < indexOf a = avlLookup o l
| o > indexOf a = avlLookup o r
avlHeight :: AVL i -> Int
avlHeight EmptyAVL = 0
avlHeight (AVL _ _ h _ _) = h
avlSize :: AVL i -> Int
avlSize EmptyAVL = 0
avlSize (AVL _ s _ _ _) = s
avlBalance :: AVL i -> AVL i
avlBalance EmptyAVL = EmptyAVL
avlBalance t@(AVL a _ _ l r)
| abs (avlHeight l avlHeight r) < 2 = t
| avlHeight l < avlHeight r = case r of
AVL a1 _ _ l1 r1 ->
let child = AVL a (findSize l l1) (findHeight l l1) l l1
in AVL a1 (findSize child r1) (findHeight child r1) child r1
| otherwise = case l of
AVL a1 _ _ l1 r1 ->
let child = AVL a (findSize r1 r) (findHeight r1 r) r1 r
in AVL a1 (findSize l1 child) (findHeight l1 child) l1 child
findSize :: AVL i -> AVL i -> Int
findSize a b = 1 + avlSize a + avlSize b
findHeight :: AVL i -> AVL i -> Int
findHeight a b = 1 + max (avlHeight a) (avlHeight b)
avlInsert :: Indexable i o v => i -> AVL i -> AVL i
avlInsert a EmptyAVL = AVL a 1 1 EmptyAVL EmptyAVL
avlInsert a (AVL a1 s h l r)
| indexOf a == indexOf a1 = AVL a s h l r
| indexOf a < indexOf a1 =
let l' = avlInsert a l
in avlBalance $ AVL a1 (s+1) (findHeight l' r) l' r
| otherwise =
let r' = avlInsert a r
in avlBalance $ AVL a1 (s+1) (findHeight l r') l r'
avlRemove :: Indexable i o v => o -> AVL i -> AVL i
avlRemove _ EmptyAVL = EmptyAVL
avlRemove o t@(AVL a _ _ EmptyAVL EmptyAVL)
| indexOf a == o = EmptyAVL
| otherwise = t
avlRemove o t@(AVL a _ _ l r)
| indexOf a == o =
case t of
AVL _ _ _ EmptyAVL _ -> case getLeft r of
(Just a',r') -> avlBalance $ AVL a' (findSize EmptyAVL r') (findHeight EmptyAVL r') EmptyAVL r'
_ -> case getRight l of
(Just a',l') -> avlBalance $ AVL a' (findSize l' r) (findHeight l' r) l' r
| o < indexOf a =
let l' = avlRemove o l
in avlBalance $ AVL a (findSize l' r) (findHeight l' r) l' r
| otherwise =
let r' = avlRemove o r
in avlBalance $ AVL a (findSize l r') (findHeight l r') l r'
getLeft :: AVL i -> (Maybe i,AVL i)
getLeft EmptyAVL = (Nothing,EmptyAVL)
getLeft (AVL a _ _ EmptyAVL EmptyAVL) = (Just a,EmptyAVL)
getLeft (AVL a _ _ EmptyAVL r) = (Just a,r)
getLeft (AVL a _ _ l r) =
case getLeft l of
(p, t2) -> (p, AVL a (findSize r t2) (findHeight r t2) t2 r)
getRight :: AVL i -> (Maybe i,AVL i)
getRight EmptyAVL = (Nothing,EmptyAVL)
getRight (AVL a _ _ EmptyAVL EmptyAVL) = (Just a,EmptyAVL)
getRight (AVL a _ _ l EmptyAVL) = (Just a,l)
getRight (AVL a _ _ l r) =
case getRight r of
(p, t2) -> (p, AVL a (findSize l t2) (findHeight l t2) l t2)
instance Functor AVL where
fmap _ EmptyAVL = EmptyAVL
fmap f (AVL a s h l r) = AVL (f a) s h (fmap f l) (fmap f r)
avlRoot :: AVL i -> i
avlRoot EmptyAVL = error "Trying to get the root of an empty AVL tree."
avlRoot (AVL a _ _ _ _) = a
avlPreorder :: AVL i -> [i]
avlPreorder EmptyAVL = []
avlPreorder (AVL a _ _ l r) = a : avlPreorder l ++ avlPreorder r
avlPostorder :: AVL i -> [i]
avlPostorder EmptyAVL = []
avlPostorder (AVL a _ _ l r) = avlPostorder l ++ avlPostorder r ++ [a]
avlInorder :: AVL i -> [i]
avlInorder EmptyAVL = []
avlInorder (AVL a _ _ l r) = avlInorder l ++ [a] ++ avlInorder r