aatree.hs

-- Arne Andersson - Balanced Search Trees Made Simple
-- Original article: http://user.it.uu.se/~arnea/ps/simp.pdf
-- Wikipedia entry: http://en.wikipedia.org/wiki/AA_tree

data AATree a = E | T a (AATree a) (AATree a) Int deriving (Show, Eq)

insert E x = T x E E 1
insert (T y left right level) x
  | x < y     = balance $ T y (insert left x) right level
  | otherwise = balance $ T y left (insert right x) level
  where balance = split . skew 

-- node names as illustration: http://en.wikipedia.org/wiki/Tree_rotation
rotateR (T q (T p a b _) c level) = T p a (T q b c level) level 
rotateL (T p a (T q b c _) level) = T q (T p a b level) c level 

-- find horizontal left edges and rotate them right
skew tree@(T _ (T _ _ _ v) _ level)
  | v == level = rotateR tree
skew tree = tree

-- if the pseudo-node is too large, split it by increasing the
-- level of every second node, by following the right path making
-- left rotations
split tree@(T _ _ (T _ _ (T _ _ _ v) _) level)
  | v == level = T x left right (level+1)
    where (T x left right _) = rotateL tree
split tree = tree

{-
funky insertion order to get the same tree as the example in the paper 
foldl insert E [4,10,2,8,12,1,3,5,9,11,13,7]
=> 
T 4
  (T 2
   (T 1 E E 1)
   (T 3 E E 1) 2)
  (T 10
   (T 8
    (T 5 E
     (T 7 E E 1) 1)
    (T 9 E E 1) 2)
   (T 12
    (T 11 E E 1)
    (T 13 E E 1) 2) 3) 3

with the 6 added
foldl insert E [4,10,2,8,12,1,3,5,9,11,13,7,6]
=>
T 4
  (T 2
   (T 1 E E 1)
   (T 3 E E 1) 2)
  (T 10
   (T 6
    (T 5 E E 1)
    (T 8
     (T 7 E E 1)
     (T 9 E E 1) 2) 2) <= both on level 2
   (T 12
    (T 11 E E 1)
    (T 13 E E 1) 2) 3) 3
-}