aatree.hs

import Data.Foldable
import qualified Data.List as L
data Tree a = Node Integer (Tree a) (Tree a) a | Leaf deriving (Show, Eq)
data Path = LeftP | RightP


instance Foldable Tree where
    foldMap f Leaf = mempty
    foldMap f (Node _ l r v) = foldMap f l `mappend` f v `mappend` foldMap f r

leftRotate :: Tree a -> Tree a
leftRotate (Node tlvl tleft (Node rlvl rleft rright rdata) tdata) =
    Node rlvl (Node tlvl tleft rleft tdata) rright rdata
leftRotate tree = tree

rightRotate (Node tlvl (Node llvl lleft lright ldata) tright tdata) =
    Node llvl lleft (Node tlvl lright tright tdata) ldata 
rightRotate tree = tree

left :: Tree a -> Tree a
left Leaf = Leaf
left (Node _ left _ _) = left

right :: Tree a -> Tree a
right Leaf = Leaf
right (Node _ _ right _) = right

level :: Tree a -> Integer
level Leaf = 0
level (Node lvl _ _ _) = lvl

isLeaf :: Tree a -> Bool
isLeaf Leaf = True
isLeaf _ = False

mapLevel :: (Integer -> Integer) -> Tree a -> Tree a
mapLevel _ Leaf = Leaf
mapLevel f (Node lvl l r d) = Node (f lvl) l r d

incLevel :: Tree a -> Tree a
incLevel = mapLevel (1+)

decLevel :: Tree a -> Tree a
decLevel = mapLevel (1-)

skew :: Tree a -> Tree a
skew tree 
    | level tree == (level . left) tree = rightRotate tree
    | otherwise = tree

split :: Tree a -> Tree a
split tree
    | level tree == (level . right . right) tree = incLevel . leftRotate $ tree
    | otherwise = tree

mapAt :: (Tree a -> Tree a) -> [Path] -> Tree a -> Tree a
mapAt f _ Leaf = f Leaf
mapAt f [] t = f t
mapAt f (RightP:xs) (Node lvl l r k) = Node lvl l (mapAt f xs r) k
mapAt f (LeftP:xs) (Node lvl l r k) = Node lvl (mapAt f xs l) r k

mapLeft :: (Tree a -> Tree a) -> Tree a -> Tree a
mapLeft f = mapAt f [LeftP]

mapRight :: (Tree a -> Tree a) -> Tree a -> Tree a
mapRight f = mapAt f [RightP]

insert :: Ord a => a -> Tree a -> Tree a
insert d tree@(Node lvl l r k)
    | d < k = split . skew . mapLeft (insert d) $ tree
    | d > k = split . skew . mapRight (insert d) $ tree
    | otherwise = tree
insert d Leaf = Node 1 Leaf Leaf d

predecessor :: Tree a -> Tree a
predecessor = goRight . left
    where goRight t
            | isLeaf t = t
            | (isLeaf . left) t && (isLeaf . right) t = t
            | otherwise = goRight . right $ t

successor :: Tree a -> Tree a
successor = goLeft . right
    where goLeft t
            | isLeaf t = t
            | (isLeaf . left) t && (isLeaf . right) t = t
            | otherwise = goLeft . left $ t

deleteRebalance :: Tree a -> Tree a
deleteRebalance Leaf = Leaf
deleteRebalance tree@(Node lvl l r k)
    -- If we have some links that jump more than one level.
    | level l < newLevel || level r < newLevel =
      splitRebalance . skewRebalance . mapLevel (return newLevel) . mapRight adjustRightLevel $ tree
    | otherwise = tree
   where 
    newLevel = min ((level . left) tree) ((level . right) tree) + 1
    adjustRightLevel = mapLevel (min newLevel)
    splitRebalance tree = foldl (flip $ mapAt split) tree $ reverse (L.tails [RightP])
    skewRebalance tree = foldl (flip $ mapAt skew) tree $ reverse (L.tails [RightP, RightP])


delete :: Ord a => a -> Tree a -> Tree a
delete d Leaf = Leaf
delete d tree@(Node lvl left right v)
    | level tree == 1 && v == d = Leaf
    | d < v = deleteRebalance . mapLeft (delete d) $ tree
    | d > v = deleteRebalance . mapRight (delete d) $ tree
    | isLeaf left = let (Node _ _ _ v) = successor tree in
                        deleteRebalance . mapRight (delete v) $ Node lvl left right v
    | otherwise = let (Node _ _ _ v) = predecessor tree in
                        deleteRebalance . mapLeft (delete v) $ Node lvl left right v

validTree :: Tree a -> Bool
validTree Leaf = True
validTree tree = (level . left) tree < level tree && 
                 (level . right) tree <= level tree && 
                 (level . right . right) tree < level tree &&
                 validTree (left tree) && validTree (right tree)