Balancing an AVL Tree haskell

Question

I am creating an AVL Tree in Haskell but I am unsure of how to balance the tree. I can add elements in but they aren't balanced. like using the addList method I add in [4,2,1,3,6,8] adds it in as:

The layout for: Root 4 (Root 2 (Root 1 Empty Empty) (Root 2 Empty Empty)) (Root 6 Empty (Root 8 Empty Empty))

should be printed as:

         4
     2       6
 1       3       8

I want to balance a tree right, but don't know how to implement it correctly, Here is my code so far, Any help on this would be amazing.

data AVLTree a = Empty | Root a (AVLTree a) (AVLTree a)
    deriving (Eq, Ord, Show)
leaf a = Root a Empty Empty

addNode :: Integral a => a -> AVLTree a -> AVLTree a
addNode a Empty = leaf a
addNode x (Root a left right)
  | x > a = Root a left (addNode x right)
  | x < a = Root a (addNode x left) right
  | otherwise = (Root a left right)
addList :: (Integral a) => [a] -> AVLTree a
addList [] = Empty
addList [n] = leaf n
addList (x:xs) = addNode x (addList xs)



search :: Integral a => AVLTree a -> a -> Bool
search Empty _ = False
search (Root a left right) x
  | x == a = True
  | x < a = search left x
  | x > a = search right x

-- balance :: AVLTree a -> AVLTree a
-- balance Empty = Empty
-- balance tree
-- |(Root r (Root left (Root left1 Empty Empty) Empty) Empty) = (Root left left1 r) -- left left case
-- |(Root r Empty (Root right Empty Empty) (Root right Empty Empty)) = (Root right r right1) -- right right case
-- |(Root r (Root left Empty (Root right Empty Empty)) Empty) = (Root r (Root right (Root left Empty Empty) Empty) Empty) -- left right case
-- |(Root r Empty (Root right (Root left Empty Empty) Empty)) = (Root r Empty (Root left Empty (Root right Empty Empty))) -- right left case

Answer 1

The implementation is missing a lot of stuff still, you should go read up a bit more on what is needed (even wikipedia has a nice description, but many other pages will show up if you do a search.)

The trickiest part of writing any (self-balancing) tree is in the balancing code.... Just creating a tree in Haskell is very easy, as you have done above, but keeping searches log(N) is a bit harder. In the case of AVL, you will have to add more code that does the following-

1. Add an int to each node that measures the height at that node. The height is measured as the distance to the furthest leaf. AVL works by verifying that the difference between the height of the left and right subnodes never differs by more than 1. It is important to note that this height needs to be added directly to the node and not calculated when needed (as some demo code on the internet does), or else even a simple insert will need to traverse the whole tree and no longer be log N.

2. Add code to walk up the node tree after an insertion (probably the best place to do this would be in insert itself, since you already have walked down the tree to the insertion point, you can just put the check after the recursion). Calculate the height (just add one to the max of the two heights below, which you just calculated in the recursive call), and check the balance factor (height of right - height of left). If things are unbalanced (ie- |balance factor| > 1) you have to....

3. Call the rotate function. Rotating is an operation with type AVLTree->AVLTree. Intuitively, it pushes the top values in the tree over to the left (or right depending on the direction of imbalance), thus restoring balance. The right node becomes the new top node, and hands its left node to the the old top, which then becomes the new left node. This is the visual-

   a                            c
  / \                          / \
 b   c             ->         a   e
    / \                      / \
   d   e                    b   d

Note that you will need to rotate either once or twice per imbalance, depending on a simple check that you will make.... If the top is imbalced (say, to the right), before you rotate the top to the left, you might need to rotate the right node to the right, if the right node has any imbalance to the left (even by 1, unlike the top, which can have -1, 0, or 1 and not need a rotate). The rotation code is the same, so you should use the same function for top and right rotates (although you might want to have separate version for rotating to the right or left).

If you run this balancing code for each insert, it will be enough.... If you ever add nodes and don't do this, you might have to repeat this process a few times across all nodes before the tree is balanced, and this will take more than log(N) computations....

As I wrote this, I noticed I barely wrote any Haskell specific information.... This info would work for any language. I suppose that is as it should be, there is no reason that the implementation should differ just because we are in Haskell. If you are having any trouble mapping to Haskell how some of this would be done, just ask in the comments and I can fill that in also.

Also note, that if you just want an object with Log(N) insert/delete, Haskell already has the Data.Sequence type. If that works for you, you won't have to write any of this yourself.

---

Edited to address comments-

The height function you mention has the problem I mentioned above- You are recursing through the whole tree to find the height of any node. Inserts will no longer have log N complexity, which negates the whole point in having a tree in the first place. (As far as I know, Haskell doesn't memoize any of these values.... If someone knows otherwise, feel free to chime in, it would simplify this code).

To elaborate a bit more, you will want to extend the definition of AVLTree like this-

data AVLTree a = Empty | Node a (AVLTree a) (AVLTree a) Int

Now the height is just obtained like this-

height (AVLTree _ _ _ height) = height

Then your balanceFactor works-

balanceFactor :: AVLTree a -> Int 
balanceFactor Empty = 0 
balanceFactor (Root a left right) = (height right) - (height left)

The downside is that you will have to add code to recalculate the height up the whole chain to the root after a change.

You will then need to write another function that checks if a rotation is needed, and applies it if needed-

rotateIfNeeded::AVLTree->AVLTree
rotateIfNeeded tree 
    | b > 1 = fullRotateLeft tree
    | b < -1 = fullRotateRight tree
    | otherwise = tree
       where b = balanceFactor tree

The functions fullRotateLeft/fullRotateRight are a wrappers around rotateLeft/rotateRight, which first check if the right/left trees need a reverse rotation first, then apply all needed rotations.

Your rotateRight/rotateLeft function will not work.... Too many variables on the left are filled in with Empty, so many cases are missing. You will want to fill most of those with variables, and move them to the appropriate place on the right.