Skip to content

Instantly share code, notes, and snippets.

@virtualsafety
Forked from gerard/AVLTree.hs
Created December 23, 2013 09:58
Show Gist options
  • Select an option

  • Save virtualsafety/8094376 to your computer and use it in GitHub Desktop.

Select an option

Save virtualsafety/8094376 to your computer and use it in GitHub Desktop.
module AVLTree where
data BT = L | N Int BT BT deriving (Show, Eq)
-- Nice, small and useful functions
empty = L
-- You could say, depth L == Nothing depth (N v L L) == Just 0, but works for
-- me better this way:
depth L = 0
depth (N _ t u) = (max (depth t) (depth u)) + 1
inorder :: BT -> [Int]
inorder L = []
inorder (N v t u) = inorder t ++ [v] ++ inorder u
left (N _ t _) = t
right (N _ _ u) = u
value (N v _ _) = v
btmin = head . inorder
-- FIXME: Could be cleaner BT -> Int using left and right of BT
balFactor :: BT -> BT -> Int
balFactor t u = (depth t) - (depth u)
-- Tricky but easy: we return a binary list with the route to the node
search :: BT -> Int -> Maybe [Int]
search L s = Nothing
search (N v t u) s
| v == s = Just []
| (search t s) /= Nothing = fmap ((:) 0) (search t s)
| (search u s) /= Nothing = fmap ((:) 1) (search u s)
| otherwise = Nothing
-- Complementary to search: get the node with the path
getelem :: BT -> [Int] -> Maybe Int
getelem L _ = Nothing
getelem (N v _ _) [] = Just v
getelem (N v t u) (x:xs)
| x == 0 = getelem t xs
| otherwise = getelem u xs
-- If you get confused (I do), check this nice picture:
-- http://en.wikipedia.org/wiki/Image:Tree_Rebalancing.gif
balanceLL (N v (N vl tl ul) u) = (N vl tl (N v ul u))
balanceLR (N v (N vl tl (N vlr tlr ulr)) u) = (N vlr (N vl tl tlr) (N v ulr u))
balanceRL (N v t (N vr (N vrl trl url) ur)) = (N vrl (N v t trl) (N vr url ur))
balanceRR (N v t (N vr tr ur)) = (N vr (N v t tr) ur)
-- Balanced insert
insert :: BT -> Int -> BT
insert L i = (N i L L)
insert (N v t u) i
| i == v = (N v t u)
| i < v && (balFactor ti u) == 2 && i < value t = balanceLL (N v ti u)
| i < v && (balFactor ti u) == 2 && i > value t = balanceLR (N v ti u)
| i > v && (balFactor t ui) == -2 && i < value u = balanceRL (N v t ui)
| i > v && (balFactor t ui) == -2 && i > value u = balanceRR (N v t ui)
| i < v = (N v ti u)
| i > v = (N v t ui)
where ti = insert t i
ui = insert u i
-- Balanced delete
delete :: BT -> Int -> BT
delete L d = L
delete (N v L L) d = if v == d then L else (N v L L)
delete (N v t L) d = if v == d then t else (N v t L)
delete (N v L u) d = if v == d then u else (N v L u)
delete (N v t u) d
| v == d = (N mu t dmin)
| v > d && abs (balFactor dt u) < 2 = (N v dt u)
| v < d && abs (balFactor t du) < 2 = (N v t du)
| v > d && (balFactor (left u) (right u)) < 0 = balanceRR (N v dt u)
| v < d && (balFactor (left t) (right t)) > 0 = balanceLL (N v t du)
| v > d = balanceRL (N v dt u)
| v < d = balanceLR (N v t du)
where dmin = delete u mu
dt = delete t d
du = delete u d
mu = btmin u
-- Test Functions
load :: BT -> [Int] -> BT
load t [] = t
load t (x:xs) = insert (load t xs) x
unload :: BT -> [Int] -> BT
unload t [] = t
unload t (x:xs) = delete (unload t xs) x
sort :: [Int] -> [Int]
sort = inorder . (load empty)
isBalanced L = True
isBalanced (N _ t u) = isBalanced t && isBalanced u && abs (balFactor t u) < 2
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment