Skip to content

Instantly share code, notes, and snippets.

Created September 13, 2013 22:26
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save anonymous/6556875 to your computer and use it in GitHub Desktop.
Save anonymous/6556875 to your computer and use it in GitHub Desktop.
AVL tree implementation in haskell
data Tree a = Nil | Node (Tree a) (Tree a) a Int deriving Show
insert :: (Ord a) => (Tree a) -> a -> (Tree a)
insert Nil x = (Node Nil Nil x 0)
insert (Node l r v _) x
| x < v = (rebalance (Node nl r v (inlineHeight nl r)))
| x > v = (rebalance (Node l nr v (inlineHeight l nr)))
where nl = insert l x
nr = insert r x
delete :: Ord a => Tree a -> a -> Tree a
delete Nil _ = Nil
delete t@(Node l r v h) x
| x < v = rebalance $ (Node (delete l x) r v (inlineHeight (delete l x) r))
| x > v = rebalance $ (Node l (delete r x) v (inlineHeight l (delete r x)))
| x == v = delete' t
delete' :: Tree a -> Tree a
delete' Nil = Nil
delete' (Node Nil Nil _ _) = Nil
delete' t@(Node l r v h)
| balance t <= 0 = rebalance $ (Node rrl (delete' rrr) rrv (inlineHeight rrl (delete' rrr)))
| otherwise = rebalance $ (Node (delete' lrl) lrr lrv (inlineHeight (delete' lrl) lrr))
where (Node lrl lrr lrv _) = leftRotate t
(Node rrl rrr rrv _) = rightRotate t
rebalance :: (Tree a) -> (Tree a)
rebalance t@(Node l r v _)
| (balance t) <= (-2) = rightRotate (Node nl r v (inlineHeight nl r))
| (balance t) >= 2 = leftRotate (Node l nr v (inlineHeight l nr))
| otherwise = t
where nl = if (balance l) > 0 then leftRotate l else l
nr = if (balance r) < 0 then rightRotate r else r
leftRotate :: (Tree a) -> (Tree a)
leftRotate (Node l r@(Node cl cr cv _) v _) = (Node nl cr cv (inlineHeight nl cr))
where nl = (Node l cl v (inlineHeight l cl))
leftRotate _ = error "can't left rotate"
rightRotate :: (Tree a) -> (Tree a)
rightRotate (Node (Node cl cr cv _) r v _) = (Node cl nr cv (inlineHeight cl nr))
where nr = (Node cr r v (inlineHeight cr r))
rightRotate t = error "can't right rotate"
balance :: (Tree a) -> Int
balance Nil = 0
balance (Node Nil Nil _ _) = 0
balance (Node l r _ _) = (height r) - (height l)
inlineHeight :: (Tree a) -> (Tree a) -> Int
inlineHeight l r = (max (height l) (height r)) + 1
height :: (Tree a) -> Int
height (Node _ _ _ h) = h
height Nil = -1
insertall :: (Ord a) => (Tree a) -> [a] -> (Tree a)
insertall t [] = t
insertall t xs = foldl insert t xs
inorder :: (Tree a) -> [(a, Int)]
inorder Nil = []
inorder (Node l r v d) = (v, d) : (inorder l ++ inorder r)
avl :: Int -> (Tree Int)
avl n = insertall Nil [1..n]
ascending = avl 10
descending = insertall Nil [10,9..1]
@tibbe
Copy link

tibbe commented Sep 14, 2013

data Tree a = Nil | Node (Tree a) (Tree a) a Int deriving Show

You want the tree to be strict in the subtrees (the two Tree a arguments). Add bang patterns (!) in front of the arguments to do that. If the three is lazy in the subtrees, insert will only do O(1) work and the rest of the computation will be performed when the result is needed e.g. when someone tries to look up an element in the tree.

You also want to mark the Int field as strict and unpacked (i.e. {-# UNPACK #-} !Int). This will improve performance and memory usage. The Int will be stored as a single machine word instead of a pointer to a heap allocated Int.

insert :: (Ord a) => (Tree a) -> a  -> (Tree a)

You want to mark every function with an Ord constraint as INLINABLE. This will make GHC create type specialized copies of the functions at the call site, improving performance quite a bit.

insert Nil x = (Node Nil Nil x 0)

Since the base case of the insert function can return a result without looking at x, this function argument cannot be passed unboxed. This likely will cost you 10% performance. Add a bang in front of x to solve this.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment