Binary Tree Operations

bst.hs

{-# OPTIONS_GHC -Wall #-}
module Tree where
import Data.List

data Tree a = Leaf
             | Node Integer (Tree a) a (Tree a)
             deriving (Show)

treeHeight :: Tree a -> Integer
treeHeight Leaf = 0
treeHeight (Node height _ _ _) = height

maxTreeHeight :: Tree a -> Tree a -> Integer
maxTreeHeight t1 t2 = maximum [(treeHeight t1), (treeHeight t2)]

balanceFactor :: Tree a -> Integer
balanceFactor Leaf = 0
balanceFactor (Node _ Leaf _ Leaf) = 0
balanceFactor (Node _ (Node lHeight _ _ _) _ (Node rHeight _ _ _)) = lHeight - rHeight
balanceFactor (Node _ Leaf _ (Node rHeight _ _ _)) = 0 - rHeight
balanceFactor (Node _ (Node lHeight _ _ _) _ Leaf) = lHeight - 0

rightRotate :: Tree a -> Tree a
rightRotate (Node _ (Node _ llTree lVal lrTree) val rTree) =
    (Node newTreeHeight llTree lVal newRTree)
    where newRTree = (Node (succ (maxTreeHeight lrTree rTree)) lrTree val rTree)
          newTreeHeight = succ (maxTreeHeight newRTree llTree)
rightRotate _ = error "rightRotate: impossible case - "

leftRotate :: Tree a -> Tree a
leftRotate (Node _ lTree val (Node _ rlTree _ rrTree)) =
    (Node newTreeHeight newLTree val rrTree)
    where newLTree = (Node (succ (maxTreeHeight lTree rlTree)) lTree val rlTree)
          newTreeHeight = succ (maxTreeHeight newLTree rrTree)
leftRotate _ = error "leftRotate: impossible case - "

balanceLeft :: Tree a -> Tree a
balanceLeft tree
    | balanceFactor tree == (-1) = leftRotate tree
    | otherwise                  = tree

balanceRight :: Tree a -> Tree a
balanceRight tree
    | balanceFactor tree == (-1) = rightRotate tree
    | otherwise                  = tree

balance :: Tree a -> Tree a
balance leaf@(Leaf) = leaf
balance tree
    | balanceFactor tree == (-2) = leftRotate (balanceLeft tree)
    | balanceFactor tree == 2    = rightRotate (balanceRight tree)
    | otherwise           = tree

balancedInsert :: Ord a => a -> Tree a -> Tree a
balancedInsert val (Leaf) = Node 0 Leaf val Leaf
balancedInsert newVal tree@(Node _ leftTree pivotVal rightTree)
    | newVal < pivotVal = newBalancedTree (balancedInsert newVal leftTree) pivotVal rightTree
    | newVal > pivotVal = newBalancedTree leftTree pivotVal (balancedInsert newVal rightTree)
    | otherwise         = tree
    where newBalancedTree lTree val rTree = balance (Node (succ (maxTreeHeight lTree rTree)) lTree val rTree)

foldTree :: Ord a => [a] -> Tree a
foldTree as =
    foldr balancedInsert Leaf as

showTree :: (Show a) => Tree a -> String
showTree Leaf = []
showTree tree = (unwords (intersperse "\n" (showTreeList tree))) ++ "\n"

showTreeList :: (Show a) => Tree a -> [String]
showTreeList (Leaf) = []
showTreeList (Node height lTree val rTree) =
        ([(show val) ++ " - " ++ (show height)] ++ (printIfNotLeaf lTree "left:") ++ (printIfNotLeaf rTree "right:"))
    where printIfNotLeaf Leaf     _ = []
          printIfNotLeaf tree label = indent [label] ++ indent (indent (showTreeList tree))
          indent = map ("  "++)