P3369 【模板】普通平衡树 avl树题解

P3369 【模板】普通平衡树

//AVL树题解 
#include<bits/stdc++.h>
using namespace std;
typedef struct AVLTreeNode{
	int key;//键值
	int height;//树的高度
	int size;//树的大小
	int cnt; //同一个键值的个数
	AVLTreeNode* left;//左儿子
	AVLTreeNode* right;//右儿子
	AVLTreeNode(int val,AVLTreeNode* l,AVLTreeNode* r):
		key(val),height(0),size(1),cnt(1),left(l),right(r){
		}//构造函数
}AVLTreeNode;
int height(AVLTreeNode* tree)
{
	if (tree!=NULL)
	{
		return tree->height;
	}
	return 0;
}
int size(AVLTreeNode* tree)
{
	if (tree!=NULL)
	{
		return tree->size;
	}
	return 0;
}
int cnt(AVLTreeNode* tree)
{
	if (tree!=NULL)
	{
		return tree->cnt;
	}
	return 0;
}
void update(AVLTreeNode* tree)//更新子树大小
{
	if (tree==NULL) return ;
	tree->size=size(tree->left)+size(tree->right)+tree->cnt;	
} 
//LL型 右旋 
AVLTreeNode* leftLeftRotation(AVLTreeNode* k2)
{
	AVLTreeNode* k1;
	k1=k2->left;
	k2->left=k1->right;
	k1->right=k2;
	k2->height=max(height(k2->left),height(k2->right))+1;
	update(k2);
	k1->height=max(height(k1->left),k2->height)+1;
	update(k1);
	return k1;
}
//RR型 左旋
AVLTreeNode* rightRightRotation(AVLTreeNode* k1)
{
	AVLTreeNode* k2;
	k2=k1->right;
	k1->right=k2->left;
	k2->left=k1;
	k1->height=max(height(k1->left),height(k1->right))+1;
	update(k1);
	k2->height=max(height(k2->right),k1->height)+1;
	update(k2);
	return k2;	
}
//LR型 双旋 先左旋再右旋
AVLTreeNode* leftRightRotation(AVLTreeNode* k3)
{
	k3->left=rightRightRotation(k3->left);
	return leftLeftRotation(k3);	
} 
//RL型 双旋 先右旋再左旋
AVLTreeNode* rightLeftRotation(AVLTreeNode* k1)
{
	k1->right=leftLeftRotation(k1->right);
	return rightRightRotation(k1);	
}
//插入
void Insert(AVLTreeNode* &tree,int key)
{
	if (tree==NULL)
	{
		tree=new AVLTreeNode(key,NULL,NULL);
	}
	else if(key<tree->key)
	{
		Insert(tree->left,key);
		if (height(tree->left)-height(tree->right)==2)
		{
			AVLTreeNode* l;
			l=tree->left;
			if (height(l->left)>height(l->right))
			{
				tree=leftLeftRotation(tree);
			}
			else
			{
				tree=leftRightRotation(tree); 
			}
		}
	}
	else
	{
		Insert(tree->right,key);
		if(height(tree->right)-height(tree->left)==2)
		{
			AVLTreeNode* r;
			r=tree->right;
			if (height(r->right)>height(r->left))
			{
				tree=rightRightRotation(tree);
			}
			else
			{
				tree=rightLeftRotation(tree); 
			}
		}
	}
	tree->height=max(height(tree->left),height(tree->right))+1;
	update(tree);
} 
AVLTreeNode* Maximum(AVLTreeNode* tree)
{
	if (tree==NULL)
	{
		return tree; 
	}
	while(tree->right!=NULL)
	{
		tree=tree->right;
	}
	return tree;
} 
AVLTreeNode* Minimum(AVLTreeNode* tree)
{
	if(tree==NULL)
	{
		return tree; 
	}
	while(tree->left!=NULL)
	{
		tree=tree->left;
	}
	return tree;
}
AVLTreeNode* Search(AVLTreeNode* tree,int key)
{
	if (tree==NULL||tree->key==key)
	{
		return tree;
	}
	if(key<tree->key)
	{
		return Search(tree->left,key); 
	}
	else
	{
		return Search(tree->right,key);
	}
}
void Remove(AVLTreeNode* &tree,int key)
{
	AVLTreeNode* z;
	if((z=Search(tree,key))!=NULL)
	{
		if (z->key<tree->key)
		{
			Remove(tree->left,key);
			if(height(tree->right)-height(tree->left)==2)
			{
				AVLTreeNode *r=tree->right;
				if (height(r->left)>height(r->right))
				{
					tree=rightLeftRotation(tree);
				}
				else
				{
					tree=rightRightRotation(tree);
				}
					
			}	
		}
		else if(z->key>tree->key)
		{
			Remove(tree->right,key);
			if(height(tree->left)-height(tree->right)==2)
			{
				AVLTreeNode *l=tree->left;
				if(height(l->right)>height(l->left))
				{
					tree=leftRightRotation(tree); 
				}
				else
				{
					tree=leftLeftRotation(tree); 
				}
			}
		}
		else
		{
			if (tree->left!=NULL&&tree->right!=NULL)
			{
				if(height(tree->left)>height(tree->right))
				{
					AVLTreeNode* Max=Maximum(tree->left);
					tree->key=Max->key;
					Remove(tree->left,Max->key); 
				}
				else
				{
					AVLTreeNode* Min=Minimum(tree->right);
					tree->key=Min->key;
					Remove(tree->right,Min->key);	
				}
			}
			else
			{
				AVLTreeNode* tmp=tree;
				tree=(tree->left!=NULL)?tree->left:tree->right;
				delete tmp; 
			}
		}
	}
	update(tree);	
}
void PreOrder(AVLTreeNode* root)
{
	if (root!=NULL)
	{
		cout<<root<<",left:"<<root->left<<",right:"<<root->right<<",height:"<<root->height<<",key:"<<root->key<<endl;
		PreOrder(root->left);
		PreOrder(root->right);		
	}
}
void InOrder(AVLTreeNode* root)
{
	if (root!=NULL)
	{
		InOrder(root->left);
		cout<<root->key<<" ";
		InOrder(root->right);
	}
}
int kth(AVLTreeNode* root,int k)
{
	if (size(root->left)+root->cnt==k) return root->key;
	else if(k<=size(root->left)) return kth(root->left,k);
	else return kth(root->right,k-size(root->left)-root->cnt);
}
int Rank(AVLTreeNode* root,int val)
{
	if (root==NULL) return 1;
	if (root->key>=val) return Rank(root->left,val);
	else return size(root->left)+root->cnt+Rank(root->right,val);	
}
int Pre(AVLTreeNode* root,int val)
{
	return kth(root,Rank(root,val)-1);
}
int suc(AVLTreeNode* root,int val)
{
	return kth(root,Rank(root,val+1));
}
int main()
{
	int n;
	cin>>n;
	AVLTreeNode* root=NULL;
	for(int i=1;i<=n;i++)
	{
		int op,x;
		cin>>op>>x;
		if (op==1)
		{
			Insert(root,x);
		}
		else if (op==2)
		{
			Remove(root,x);
		}
		else if(op==3)
		{
			cout<<Rank(root,x)<<endl;
		}
		else if (op==4)
		{
			cout<<kth(root,x)<<endl;
		}
		else if(op==5)
		{
			cout<<Pre(root,x)<<endl;
		}
		else
		{
			cout<<suc(root,x)<<endl;
		}
		
	}
}