Created
September 13, 2013 22:26
-
-
Save anonymous/6556875 to your computer and use it in GitHub Desktop.
AVL tree implementation in haskell
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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. TheInt
will be stored as a single machine word instead of a pointer to a heap allocatedInt
.You want to mark every function with an
Ord
constraint asINLINABLE
. This will make GHC create type specialized copies of the functions at the call site, improving performance quite a bit.Since the base case of the
insert
function can return a result without looking atx
, this function argument cannot be passed unboxed. This likely will cost you 10% performance. Add a bang in front ofx
to solve this.