AVL Tree

avl_tree.hpp

#include "bst_tree.hpp"

template<typename KeyType, typename Comparator = std::less<KeyType> >
class avl_tree : public binary_search_tree<KeyType, Comparator>{
protected:
    class avl_node : public binary_search_tree<KeyType, Comparator>::bst_node{
    public:
        avl_node(KeyType key):
            binary_search_tree<KeyType, Comparator>::bst_node(key),height(1){}
        int height;
    };
private:
    avl_node * right_rotate(avl_node *y);
    avl_node * left_rotate(avl_node *y);
public:
    typename binary_search_tree<KeyType, Comparator>::iterator insert(KeyType key);
};

template <typename KeyType, typename Comparator>
typename avl_tree<KeyType, Comparator>::avl_node *avl_tree<KeyType, 
                            Comparator>::right_rotate(avl_node *y)
{
    avl_node *x,*t2;
    x = (avl_node*)y->left;	//Has to  be a genuine node
    bool yRightChildExists = y != NULL && y->right != NULL 
                            && y->num_child > y->right->num_child;
    bool xLeftChildExists = x!=NULL && x->left != NULL 
                            && x->num_child > x->left->num_child;
    if (x->right != NULL && x->right->num_child < x->num_child)
    {
        t2 = (avl_node*)x->right;
        y->left = t2;	   //Genuine Child Node
        y->num_child = 1 + (t2 != NULL ? t2->num_child : 0  +
                        yRightChildExists ? y->right->num_child : 0);
        y->height = 1 + std::max( t2!=NULL? t2->height : 0 ,
                        yRightChildExists ? ((avl_node*)y->right)->height : 0);
    }
    else
    {			            //Threaded Link
        y->left = x;
        y->num_child = 1 + (yRightChildExists ? y->right->num_child : 0);
        y->height = 1 + (yRightChildExists ? ((avl_node*)y->right)->height : 0);
    }
    x->right = y;
    x->num_child = 1 + (y != NULL ? y->num_child :0 
                    + xLeftChildExists ? x->left->num_child: 0);
    x->height = 1 + std::max(y != NULL ? y->height : 0,
                    xLeftChildExists ? ((avl_node*)x->left)->height:0);
    return x;
}

template <typename KeyType, typename Comparator>
typename avl_tree<KeyType, Comparator>::avl_node 
                  *avl_tree<KeyType, Comparator>::left_rotate(avl_node *y)
{
    avl_node *x, *t2;
    x = (avl_node*)y->right;   //Assuming to  be a genuine node passed
    bool leftChildExists = (y != NULL && y->left != NULL
                            && y->num_child > y->left->num_child);
    bool xRightChildExists = x != NULL && x->right != NULL
                            && x->num_child > x->right->num_child;
    if (x->left != NULL && x->left->num_child < x->num_child)
    {
        t2 = (avl_node*)x->left;
        y->right = t2;	//Genuine Child Node
        y->num_child = 1 + ((t2 != NULL)? t2->num_child : 0 +
                       leftChildExists ? y->left->num_child:0);
        y->height = 1 + std::max( t2!= NULL ? t2->height : 0,
                      leftChildExists? ((avl_node*)y->left)->height :0);
    }
    else
    {		//Threaded Link
        y->right = x;
        y->num_child = 1 + (leftChildExists ? y->left->num_child:0);
        y->height = 1 + (leftChildExists ? ((avl_node*)y->left)->height : 0);
    }
    x->left = y;
    x->num_child = 1 + (y != NULL ? y->num_child : 0  +
                                   xRightChildExists ? x->right->num_child :0);
    x->height = 1 + std::max( (y != NULL) ? y->height : 0 ,
                            xRightChildExists ? ((avl_node*)x->right)->height : 0);
    return x;
}

template <typename KeyType, typename Comparator>
typename binary_search_tree<KeyType, Comparator>::iterator
avl_tree<KeyType, Comparator>::insert(KeyType key) {
    std::stack<typename binary_search_tree<KeyType, Comparator>::bst_node*> parents;
    avl_node * inserted_node = new avl_node(key);
    typename binary_search_tree<KeyType, Comparator>::iterator it =
                                this->__insert(inserted_node, parents);
    if (it == this->end()){
        delete inserted_node;
        return it;
    }
    avl_node *node_ptr = NULL;
    avl_node *avl_inserted = NULL;
    while ( !parents.empty())
        {
            node_ptr = (avl_node*)parents.top();
            parents.pop();
            node_ptr->num_child++;
            int leftHeight = 0;
            if ( node_ptr->left != NULL
                 && node_ptr->num_child > node_ptr->left->num_child){
                leftHeight = ((avl_node*)node_ptr->left)->height;
            }
            int rightHeight = 0;
            if ( node_ptr->right != NULL
                 && node_ptr->num_child > node_ptr->right->num_child){
                rightHeight = ((avl_node*)node_ptr->right)->height;
            }
            node_ptr->height = 1 + std::max(leftHeight, rightHeight) ;
            int balance = leftHeight - rightHeight;
            if (balance > 1)
            {
                if (this->isLessThan(node_ptr->key, node_ptr->left->key))
                {
                    avl_inserted = right_rotate(node_ptr);
                }
                else
                {
                    node_ptr->left = left_rotate((avl_node*)node_ptr->left);
                    avl_inserted = right_rotate(node_ptr);
                }
            }
            else if (balance < -1)
            {
                if (this->isLessThan(node_ptr->key, node_ptr->right->key ))
                {
                    avl_inserted = left_rotate(node_ptr);
                }
                else
                {
                    node_ptr->right = right_rotate((avl_node*)node_ptr->right);
                    avl_inserted = left_rotate(node_ptr);
                }
            }
            if (avl_inserted){
                if (!parents.empty()){
                    typename binary_search_tree<KeyType, Comparator>::bst_node
                            * parent_ptr = parents.top();
                    if (avl_inserted->key < parent_ptr->key)
                        parent_ptr->left = avl_inserted;
                    else
                        parent_ptr->right = avl_inserted;
                }
                else
                {
                    this->root = avl_inserted;
                }
                avl_inserted = NULL;
            }
        }
        return typename binary_search_tree<KeyType, Comparator>::iterator(inserted_node);
}