AVL TREE insert， delete，rotation，algorithm

#include "AVL.h"
#include <iostream>
#include <math.h>

using namespace std;

//////////////////////////////Queue//////////////////////////////////

Queue::~Queue()
{
	QueueNode *curr = front, *next;

	while(curr)
	{
		next = curr->next;
		delete curr;
		curr = next;
	}
}

void Queue::enqueue(AVLNode *data, bool isLeading)
{
	QueueNode *newNode = new QueueNode(data, isLeading);
	if(!front)
		//The queue is empty
		front = newNode;		
	
	else
		//Insert as the new rear
		rear->next = newNode;
	rear = newNode;	
}

QueueNode* Queue::dequeue()
{
	if(!front)
	{
		//The queue is empty
		return NULL;
	}
	else
	{
		//Return the first node
		QueueNode *tmp = front;
		front = front->next;
		return tmp;
	}
}


///////////////////////////////////AVL tree///////////////////////////////////
void AVLTree::destroy(AVLNode *subRoot)
{
	if(!subRoot)
		return;
	destroy(subRoot->left);
	cout<<"Node "<<subRoot->key<<" deleted."<<endl;
	destroy(subRoot->right);
	delete subRoot;
}

int AVLTree::rGetHeight(AVLNode *subRoot)
{
	if(!subRoot)
		return 0;
	if(subRoot->balance == -1)
		//subRoot's left subtree is heavier
		return rGetHeight(subRoot->left) + 1;
	else
		return rGetHeight(subRoot->right) + 1;
}

void AVLTree::printTree()
{
	//Breadth first tree traversal
	if(!root) cout<<"Empty tree."<<endl;

	int treeHeight = this->getHeight();
	Queue printQueue;
	printQueue.enqueue(root, true);
	QueueNode *curr;	
	int space = pow(2,treeHeight+2)-1;
	//The spaces between tree nodes when printing the tree, nodes of a larger depth have larger space

	while(treeHeight >= -1)
	{
		curr  = printQueue.dequeue();
		//Print the space
		if(curr->isLeading)
		{
			treeHeight--;
			cout<<endl;
			space /= 4;
			for (int i=0; i<space; i++)
				cout<<" ";
			space = space*2 +1;
		}
		else
		{
			for (int i=0; i<space; i++)
				cout<<" ";
		}
	
		
		if(curr->data)
		{
			//curr is a tree node
			cout<<curr->data->key;
			printQueue.enqueue(curr->data->left, curr->isLeading);
			printQueue.enqueue(curr->data->right);
		}
		else
		{
			//curr is a hole
			cout<<" ";
			printQueue.enqueue(NULL, curr->isLeading);
			printQueue.enqueue(NULL);
		}		
	}
}


AVLNode* AVLTree::insert(int key)
{
	if(root)
	{
		//The tree is not empty, resursively insert the key
		bool heightChanged;
		AVLNode *newNode = rInsert(key, root, heightChanged);		
		return newNode;
	}
	else
	{
		//The tree is empty, insert the new node as the tree root
		return root = new AVLNode(key);
	}	
}

AVLNode* AVLTree::rInsert(int key, AVLNode* &subRoot, bool &heightChanged)
//heightChanged is true is the subtree leaded by subRoot is taller after insertion
{
	if(key == subRoot->key)
	{
		cout<<"No duplicate key allowed."<<endl;
		return NULL;
	}
	else if (key < subRoot->key)		
	{		
		if(subRoot->left)
		{
			//Recursively insert into subRoot's left subtree
			AVLNode *newNode = rInsert(key, subRoot->left, heightChanged);
			if(!heightChanged)
				return newNode;
			subRoot->balance--;
			switch(subRoot->balance)
			{
			case -1:
				break;
			case 0:
				heightChanged = false;
				break;
			case -2:
				//subRoot is unbalanced, and it is left heavy
				if(subRoot->left->balance == -1)
					//subRoot's left child is also left heavy, case 1
					balance1(subRoot);
				else
					//subRoot's left child is right heavy, case 2
					//Note: subRoot->left->balance can only be 1 or -1
					balance2(subRoot);
				heightChanged = false;
				break;
			}
			return newNode;			
		}
		else			
		{
			//Insert as subRoot's left child
			subRoot->balance--;
			//The balance of subRoot after insertion can only be -1 or 0
			if(subRoot->balance == -1)
				heightChanged = true;
			else
				heightChanged = false;
			return subRoot->left = new AVLNode(key);
		}
	}
	else
	{
		if(subRoot->right)
		{
			//Recursively insert into subRoot's right subtree
			AVLNode *newNode = rInsert(key, subRoot->right, heightChanged);
			if(!heightChanged)
				return newNode;
			subRoot->balance++;
			switch(subRoot->balance)
			{
			case 1:
				break;
			case 0:
				heightChanged = false;
				break;
			case 2:
				//subRoot is unbalanced, and it is right heavy
				if(subRoot->right->balance == -1)
					//subRoot's right child is left heavy, case 3
					balance3(subRoot);
				else
					//subRoot's right child is right heavy, case 4
					//Note: subRoot->left->balance can only be 1 or -1
					balance4(subRoot);
				heightChanged = false;
				break;
			}
			return newNode;			
		}
		else			
		{
			//Insert as subRoot's right child
			subRoot->balance++;
			//The balance of subRoot after insertion can only be 1 or 0
			if(subRoot->balance == 1)
				heightChanged = true;
			else
				heightChanged = false;
			return subRoot->right = new AVLNode(key);
		}
	}		
}

AVLNode* AVLTree::remove(int key)
{	// To judge the balance if it is changed
	bool heightChanged;
	if(!root){
		cout<<"The tree is empty!"<<endl;
	}
	else{
		rRemove(key,root,heightChanged);
	}
}


AVLNode *AVLTree::rRemove(int key,AVLNode* &subroot, bool &heightChanged){
	// When the key equals the current node key
	if(key==subroot->key){
		 //Left child is empty
		if ( subroot->left == NULL && subroot->right != NULL  )      
        {
        	//let temp point to subroot
        	temp=subroot;
            heightChanged=true;
            //change subroot value to the right child
            subroot=subroot->right;
            //pick the temp from the tree
            return temp;
        }
        //Right child is empty
        else if ( subroot->right == NULL && subroot->left != NULL)  
        {
        	//let temp point to subroot
        	temp=subroot;
            heightChanged=true;
             //change subroot value to the left child
            subroot=subroot->left;
            //pick the temp from the tree
            return temp;
        }
        else if(subroot->right == NULL && subroot->left == NULL){
        	temp=subroot;
  	        heightChanged=true;
  	        subroot=NULL;
       	    return temp;
        }
        // Both left and right tree
        else                        
        {
            temp = findMin( subroot->right, heightChanged);
            temp->left=subroot->left;
            temp->right=subroot->right;
            //change the balance of new subroot as the old one
            temp->balance=subroot->balance;
            subroot=temp;
			if(heightChanged){
				subroot->balance--;
			}
            if (subroot->balance==-2){
            	switch(subroot->left->balance){
	            	case -1: balance1(subroot);break;  //left-left
	            	case 1:  balance2(subroot);break;  //left-right
	            	default: break;
	            }
            	
            }else if(subroot->balance == 2){
            		switch(subroot->right->balance){
	            	case 1:  balance3(subroot);        //right-left
	            	case -1: balance4(subroot);        //right-right
	            	default: break;
	            }
            	
            }
            
        }
        //delete temp;
        //get into the left child of subroot and judge if there is the key
	}else if(key < subroot->key && subroot->left != NULL){
		 temp=rRemove( key, subroot->left, heightChanged);
		 if(heightChanged){
		 		subroot->balance++;
 			switch(subroot->balance){
			 	case 0:  heightChanged=true;break;
			 	case 1:  heightChanged=false;break;
			 	case 2:  if(subroot->right->balance == -1){
							//subRoot's right child is left heavy, case 3
					        balance3(subroot);
					        heightChanged=true;break;
				         }else if(subroot->right->balance == 0){
							//subRoot's right child is right heavy, case 4
							//Note: subRoot->left->balance can only be 1 or -1
				           	balance6(subroot);
				           	heightChanged=false;break;   //balance 5 and 6 heighted equals false
				         }else{
         					balance4(subroot);
         					heightChanged=true;break;
         				}
			 }
			 return subroot;
 		}
       
       //get into right child and judge if there is the key
	}else if(key>subroot->key && subroot->right != NULL){
 		 temp = rRemove( key, subroot->right, heightChanged);
		 if(heightChanged){  
 			subroot->balance--;
 			switch(subroot->balance){
	 			case 0:  heightChanged=true;break;
	 			case -1: heightChanged=false;break;
 				case -2: if(subroot->left->balance == -1){
							//subRoot's right child is left heavy, case 3
					        balance1(subroot);
					        heightChanged=true;break;
				         }else if(subroot->left->balance == 0){
							//subRoot's right child is right heavy, case 4
							//Note: subRoot->left->balance can only be 1 or -1
				           	balance5(subroot);
				           	heightChanged=false;break;
				         }else{
         					balance2(subroot);
         					heightChanged=true;break;
         				}
               }
                return subroot;	
          }
   }else{
   		cout<<"这个节点不存在!"<<endl;
   		return NULL;
   }
}

//Find the smallest one from the right tree
AVLNode * AVLTree::findMin( AVLNode * &subroot, bool &heightChanged)
/* Begin of findMin */
{
	tmp = subroot;
    if (subroot == NULL){
        return NULL;
    }else if (subroot -> left == NULL && subroot -> right != NULL){
    	//change the pointer 
        subroot=subroot->right;
        heightChanged=true;
        return tmp;
    }
    else if(subroot -> left == NULL && subroot -> right == NULL){
    	heightChanged = true;
    	subroot = NULL;
    	return tmp;
    }else
        tmp = findMin(subroot -> left, heightChanged);
        if(heightChanged){
        	//right is heavier
        	subroot->balance++;
        }
        switch(subroot->balance	){
        	case 0: heightChanged=true;break;
        	case 1: heightChanged=false;break;
        	case 2: if(subroot->right->balance == -1){
							//subRoot's right child is left heavy, case 3
					        balance3(subroot);
					        heightChanged=true;break;
        	                }
				         else if(subroot->right->balance == 1){
							//subRoot's right child is right heavy, case 4
							//Note: subRoot->left->balance can only be 1 or -1
				           	balance4(subroot);
				           	heightChanged=true;break;
				         }else{
				         	//subroot->right->balance == 0
         					balance6(subroot);
         					heightChanged=false;break;  //special case
         				} 	
        }
        return tmp;
        
}/* End of findMin() */



void AVLTree::balance1(AVLNode* &subRoot)
{
	//For a diagram, see slide 13, lecture 10
	AVLNode *k1 = subRoot->left, *k2 = subRoot;

	//fix k2
	k2->left = k1->right;
	k2->balance = 0;

	//fix k1
	k1->right = k2;
    k1->balance = 0;	
    subRoot = k1;	
}

void AVLTree::balance2(AVLNode* &subRoot)
{
	//For a diagram, see slide 7, lecture 11
	AVLNode *k3 = subRoot, *k1 = subRoot->left, *k2 = k1->right;

	//fix k1
	k1->right = k2->left;
	if(k2->balance == 1)	
		//if k2 is right heavy, then the height of subtree B is 1 smaller than the other subtrees, therefore k1 is left heavy
		k1->balance = -1;
	else
		k1->balance = 0;

	//fix k3
	k3->left = k2->right;
	if(k2->balance == -1)
		//if k2 is left heavy, then the height of subtree C is 1 smaller than the other subtrees, therefore k3 is right heavy
		k3->balance = 1;
	else
		k3->balance = 0;

	//fix k2
	k2->left = k1;
	k2->right = k3;
	k2->balance = 0;

	subRoot = k2;
}

void AVLTree::balance3(AVLNode* &subRoot)
{
	//For a diagram, see slide 8, lecture 11
	AVLNode *k1 = subRoot, *k3 = subRoot->right, *k2 = k3->left;

	//fix k1
	k1->right = k2->left;
	if(k2->balance == 1)
		//if k2 is right heavy, then the height of subtree B is 1 smaller than the other subtrees, therefore k1 is left heavy
		k1->balance = -1;
	else
		k1->balance = 0;

	//fix k3
	k3->left = k2->right;
	if(k2->balance == -1)
		//if k2 is left heavy, then the height of subtree C is 1 smaller than the other subtrees, therefore k3 is right heavy
		k3->balance = 1;
	else
		k3->balance = 0;

	//fix k2
	k2->left = k1;
	k2->right = k3;
	k2->balance = 0;

	subRoot = k2;
}

void AVLTree::balance4(AVLNode* &subRoot)
{
	//For the diagram, see slide 15, lecture 10
	AVLNode *k1 = subRoot, *k2 = subRoot->right;
	
	//fix k1
	k1->right = k2->left;
	k1->balance = 0;

	//fix k2
	k2->left = k1;
 	k2->balance = 0;
	subRoot = k2;
}

void AVLTree::balance5(AVLNode* &subRoot)
{
	//For a diagram, see slide 13, lecture 10
	AVLNode *k1 = subRoot->left, *k2 = subRoot;

	//fix k2
	k2->left = k1->right;

	//fix k1
	k1->right = k2;
	k1->balance=1;  //After the rotation k1 balance is 1, k2 is -1
	k2->balance=-1;
	subRoot = k1;	
}

void AVLTree::balance6(AVLNode* &subRoot)
{
	//For the diagram, see slide 15, lecture 10
	AVLNode *k1 = subRoot, *k2 = subRoot->right;
	
	//fix k1
	k1->right = k2->left;
	//fix k2
	k2->left = k1;
	k1->balance=1;    //After the rotation k1 balance is 1, k2 is -1
	k2->balance=-1;
 	subRoot = k2;
}

int main()
{
	AVLTree tree;
	char loopFlag = '1';
	int i = 8;
	int removeElement;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 5;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 4;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 2;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 3;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 0;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 15;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 10;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 11;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 17;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	i = 9;
	cout<<"Inserting "<<i<<endl;
	tree.insert(i);
	tree.printTree();
	do
	{
		cout<<"Please input element you want to remove:"<<endl;
	    cin>>removeElement;
	    tree.remove(removeElement);
	    tree.printTree();
	    cout<<"\n\n";
	    cout<<"想继续删除吗？0：退出/1：继续...."<<endl;
	    cin>>loopFlag;
	}while((loopFlag != '0'));

	return 0;
}