AA Tree in Haskell

Here’s an implementation of an AA tree, which is a simplified balanced binary tree algorithm. I’m just trying to sharpen my very rusty Haskell programming skillz. I wish the popular match macro for Scheme had guards like Haskell. Instead, the match macro has an uber-powerful fail continuation that is overkill for most situations. Anyway, this code isn’t so great, but it seems to work.

data AATree a = Node { level :: Int, 
                       left :: AATree a, 
                       value :: a, 
                       right :: AATree a }
              | Nil deriving (Show, Eq)

skew :: AATree a -> AATree a
skew (Node n (Node ln ll lv lr) v r) 
    | ln == n = Node ln ll lv (Node n lr v r)
skew t = t

split :: AATree a -> AATree a
split (Node tn a tv (Node rn b rv x@(Node xn _ _ _))) 
    | tn == xn = Node (rn+1) (Node tn a tv b) rv x
split t = t

insert :: Ord a => AATree a -> a -> AATree a
insert (Node n l v r) x | x < v = fixup (Node n (insert l x) v r)
                        | x > v = fixup (Node n l v (insert r x))
                        where fixup = split . skew 
insert Nil x = Node 1 Nil x Nil 

tpred :: AATree a -> a
tpred (Node _ _ v Nil) = v
tpred (Node _ _ v r)   = tpred r 

tsucc :: AATree a -> a
tsucc (Node _ Nil v _) = v 
tsucc (Node _ l v _)   = tsucc l 

lvl :: AATree a -> [Int]
lvl Nil = [] 
lvl (Node n _ _ _) = [n]

declvl :: AATree a -> AATree a
declvl t@(Node n Nil v Nil) = t 
declvl (Node n l v Nil) | s < n = Node s l v Nil
                        where s = level l + 1 
declvl (Node n l v (Node rn rl rv rr)) 
    | s < n && s < rn = Node s l v (Node s rl rv rr)
    where s = minimum (lvl l ++ [rn]) + 1 
declvl t = t

remove :: Ord a => AATree a -> a -> AATree a
remove (Node n l v r) x | x > v = Node n l v (remove r x)
                        | x < v = Node n (remove l x) v r
remove (Node n Nil v Nil) x | x == v = Nil 
remove (Node n Nil v r) x | x == v = Node n Nil s (remove r s)
                          where s = tsucc r
remove (Node n l v r) x | x == v = Node n (remove l p) p r
                        where p = tpred l 
remove t x = t

delete :: Ord a => AATree a -> a -> AATree a
delete Nil x = Nil 
delete t x 
    | a == Nil = Nil 
    | right a == Nil = let Node n l v r = split a in Node n l v (split r)
    | otherwise = let Node n l v r = a
                      Node rn rl rv rr = skew r 
                      Node b c d e = split $ Node n l v (Node rn rl rv (skew rr))
                  in Node b c d (split e) 
    where a = skew (declvl (remove t x))
Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s