{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}
{-
This module is part of Chatty.
Copyleft (c) 2014 Marvin Cohrs
All wrongs reversed. Sharing is an act of love, not crime.
Please share Chatty with everyone you like.
Chatty is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
Chatty is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License
along with Chatty. If not, see <http://www.gnu.org/licenses/>.
-}
-- | Provides an AVL tree.
module Data.Chatty.AVL (avlMax,avlMin,avlLookup,avlHeight,avlSize,avlInsert,avlRemove,AVL (EmptyAVL,AVL),avlRoot,avlPreorder,avlPostorder,avlInorder) where
import Data.Chatty.BST
import Data.Chatty.None
instance Indexable i o v => AnyBST AVL i o v where
anyBstMax = avlMax
anyBstMin = avlMin
anyBstLookup = avlLookup
anyBstEmpty = EmptyAVL
anyBstInsert = avlInsert
anyBstRemove = avlRemove
anyBstHead = avlHead
anyBstInorder = avlInorder
instance None (AVL a) where
none = EmptyAVL
-- | An AVL tree.
data AVL a = EmptyAVL | AVL a Int Int !(AVL a) !(AVL a)
-- | Get the greatest element.
avlMax :: AVL i -> Maybe i
avlMax EmptyAVL = Nothing
avlMax (AVL a _ _ _ EmptyAVL) = Just a
avlMax (AVL _ _ _ _ r) = avlMax r
-- | Get the least element.
avlMin :: AVL i -> Maybe i
avlMin EmptyAVL = Nothing
avlMin (AVL a _ _ EmptyAVL _) = Just a
avlMin (AVL _ _ _ l _) = avlMin l
-- | Lookup a given key.
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
-- | Get the height of the tree.
avlHeight :: AVL i -> Int
avlHeight EmptyAVL = 0
avlHeight (AVL _ _ h _ _) = h
-- | Get the size of the tree.
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)
-- | Insert into the tree.
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'
-- | Remove from the tree.
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)
-- | Get the root of the tree.
avlRoot :: AVL i -> i
avlRoot EmptyAVL = error "Trying to get the root of an empty AVL tree."
avlRoot (AVL a _ _ _ _) = a
-- | Get the root of the tree (safely)
avlHead :: AVL i -> Maybe i
avlHead EmptyAVL = Nothing
avlHead t = Just $ avlRoot t
-- | Traverse the tree, order (head, left, right)
avlPreorder :: AVL i -> [i]
avlPreorder EmptyAVL = []
avlPreorder (AVL a _ _ l r) = a : avlPreorder l ++ avlPreorder r
-- | Traverse the tree, order (left, right, head)
avlPostorder :: AVL i -> [i]
avlPostorder EmptyAVL = []
avlPostorder (AVL a _ _ l r) = avlPostorder l ++ avlPostorder r ++ [a]
-- | Traverse the tree, order (left, head, right)
avlInorder :: AVL i -> [i]
avlInorder EmptyAVL = []
avlInorder (AVL a _ _ l r) = avlInorder l ++ [a] ++ avlInorder r