balancing an AVL tree (C++)

gregghz picture gregghz · Nov 18, 2010 · Viewed 38.4k times · Source

I'm having the hardest time trying to figure out how to balance an AVL tree for my class. I've got it inserting with this:

Node* Tree::insert(int d)
{
    cout << "base insert\t" << d << endl;
    if (head == NULL)
        return (head = new Node(d));
    else
        return insert(head, d);
}

Node* Tree::insert(Node*& current, int d)
{
    cout << "insert\t" << d << endl;
    if (current == NULL)
        current = new Node(d);
    else if (d < current->data) {
        insert(current->lchild, d);
        if (height(current->lchild) - height(current->rchild)) {
            if (d < current->lchild->getData())
                rotateLeftOnce(current);
            else
                rotateLeftTwice(current);
        }
    }
    else if (d > current->getData()) {
        insert(current->rchild, d);
        if (height(current->rchild) - height(current->lchild)) {
            if (d > current->rchild->getData())
                rotateRightOnce(current);
            else
                rotateRightTwice(current);
        }
    }

    return current;
}

My plan was to have the calls to balance() check to see if the tree needs balancing and then balance as needed. The trouble is, I can't even figure out how to traverse the tree to find the correct unbalanced node. I know how to traverse the tree recursively, but I can't seem to translate that algorithm into finding the lowest unbalanced node. I'm also having trouble writing an iterative algorithm. Any help would be appreciated. :)

Answer

Carlos picture Carlos · Nov 18, 2010

You can measure the height of a branch at a given point to calculate the unbalance

(remember a difference in height (levels) >= 2 means your tree is not balanced)

int Tree::Height(TreeNode *node){
     int left, right;

     if(node==NULL)
         return 0;
     left = Height(node->left);
     right = Height(node->right);
  if(left > right)
            return left+1;
         else
            return right+1;
} 

Depending on the unevenness then you can rotate as necessary

void Tree::rotateLeftOnce(TreeNode*& node){
     TreeNode *otherNode;

     otherNode = node->left;
     node->left = otherNode->right;
     otherNode->right = node;
     node = otherNode;
}


void Tree::rotateLeftTwice(TreeNode*& node){
     rotateRightOnce(node->left);
     rotateLeftOnce(node);
}


void Tree::rotateRightOnce(TreeNode*& node){
     TreeNode *otherNode;

     otherNode = node->right;
     node->right = otherNode->left;
     otherNode->left = node;
     node = otherNode;
}


void Tree::rotateRightTwice(TreeNode*& node){
     rotateLeftOnce(node->right);
     rotateRightOnce(node);
}

Now that we know how to rotate, lets say you want to insert a value in the tree... First we check whether the tree is empty or not

TreeNode* Tree::insert(int d){
     if(isEmpty()){
         return (root = new TreeNode(d));  //Is empty when root = null
     }
     else
         return insert(root, d);           //step-into the tree and place "d"
}

When the tree is not empty we use recursion to traverse the tree and get to where is needed

TreeNode* Tree::insert(TreeNode*& node, int d_IN){
     if(node == NULL)  // (1) If we are at the end of the tree place the value
         node = new TreeNode(d_IN);
     else if(d_IN < node->d_stored){  //(2) otherwise go left if smaller
         insert(node->left, d_IN);    
         if(Height(node->left) - Height(node->right) == 2){
            if(d_IN < node->left->d_stored)
                rotateLeftOnce(node);
            else
                rotateLeftTwice(node);
         }
     }
     else if(d_IN > node->d_stored){ // (3) otherwise go right if bigger
        insert(node->right, d_IN);
        if(Height(node->right) - Height(node->left) == 2){
            if(d_IN > node->right->d_stored)
                rotateRightOnce(node);
            else
                rotateRightTwice(node);
        }
     }
     return node;
}

You should always check for balance (and do rotations if necessary) when modifying the tree, no point waiting until the end when the tree is a mess to balance it. That just complicates things...


UPDATE

There is a mistake in your implementation, in the code below you are not checking correctly whether the tree is unbalanced. You need to check whether the height is equals to 2 (therefore unbalance). As a result the code bellow...

if (height(current->lchild) - height(current->rchild)) { ...

if (height(current->rchild) - height(current->lchild)) {...

Should become...

if (height(current->lchild) - height(current->rchild) == 2) { ...

if (height(current->rchild) - height(current->lchild) == 2) {...

Some Resources