红黑树,求调
查看原帖
红黑树,求调
513746
Ay2024_4202yA楼主2025/6/14 18:59
#include<bits/stdc++.h>
using namespace std;
enum Color{RED,BLACK};
template<typename T>
class RBTree{
private:
	struct node{
		T key;
		int cnt,siz;
		Color color;
		node *left,*right,*parent;
		node(T k,Color c,node* nil):
			key(k),color(c),cnt(1),siz(1),
			left(nil),right(nil),parent(nil){}
	};
	node *root;
	node *nil;
	node* createnil(){
		node* Node=new node(T(),BLACK,nullptr);
		Node->left=Node->right=Node->parent=Node;
		Node->cnt=Node->siz=0;
		return Node;
	}
	void updatesiz(node* x){
		if(x==nil)return;
		x->siz=x->left->siz+x->right->siz+x->cnt;
	}
	void leftrotate(node* x){
		node *y=x->right;
		x->right=y->left;
		if(y->left!=nil){
			y->left->parent=x;
		}
		y->parent=x->parent;
		if(x->parent==nil){
			root=y;
		}else if(x==x->parent->left){
			x->parent->left=y;
		}else{
			x->parent->right=y;
		}
		y->left=x;
		x->parent=y;
		updatesiz(x);
		updatesiz(y);
	}
	void rightrotate(node *y){
		node *x=y->left;
		y->left=x->right;
		if(x->right!=nil){
			x->right->parent=y;
		}
		x->parent=y->parent;
		if(y->parent==nil){
			root=x;
		}else if(y==y->parent->left){
			y->parent->left=x;
		}else{
			y->parent->right=x;
		}
		x->right=y;
		y->parent=x;
		updatesiz(y);
		updatesiz(x);
	}
	void insertfixup(node *z){
		while(z->parent->color==RED){
			if(z->parent==z->parent->parent->left){
				node *y=z->parent->parent->right;
				if(y->color==RED){
					z->parent->color=BLACK;
					y->color=BLACK;
					z->parent->parent->color=RED;
					z=z->parent->parent;
				}else{
					if(z==z->parent->right){
						z=z->parent;
						leftrotate(z);
					}
					z->parent->color=BLACK;
					z->parent->parent->color=RED;
					rightrotate(z->parent->parent);
				}
			}else{
				node *y=z->parent->parent->left;
				if(y->color==RED){
					z->parent->color=BLACK;
					y->color=BLACK;
					z->parent->parent->color=RED;
					z=z->parent->parent;
				}else{
					if(z==z->parent->left){
						z=z->parent;
						rightrotate(z);
					}
					z->parent->color=BLACK;
					z->parent->parent->color=RED;
					leftrotate(z->parent->parent);
				}
			}
		}
		root->color=BLACK;
	}
	node* minimum(node *x)const{
		while(x->left!=nil){
			x=x->left;
		}
		return x;
	}
	node* maximum(node *x)const{
		while(x->right!=nil){
			x=x->right;
		}
		return x;
	}
	void transplant(node *u,node *v){
		if(u->parent==nil){
			root=v;
		}else if(u==u->parent->left){
			u->parent->left=v;
		}else{
			u->parent->right=v;
		}
		if(v!=nil)v->parent=u->parent;
	}
	void deletefixup(node *x){
		while(x!=root&&x->color==BLACK){
			if(x==x->parent->left){
				node *w=x->parent->right;
				if(w->color==RED){
					w->color=BLACK;
					x->parent->color=RED;
					leftrotate(x->parent);
					w=x->parent->right;
				}
				if(w->left->color==BLACK&&w->right->color==BLACK){
					w->color=RED;
					x=x->parent;
				}else{
					if(w->right->color==BLACK){
						w->left->color=BLACK;
						w->color=RED;
						rightrotate(w);
						w=x->parent->right;
					}
					w->color=x->parent->color;
					x->parent->color=BLACK;
					w->right->color=BLACK;
					leftrotate(x->parent);
					x=root;
				}
			}else{
				node *w=x->parent->left;
				if(w->color==RED){
					w->color=BLACK;
					x->parent->color=RED;
					rightrotate(x->parent);
					w=x->parent->left;
				}
				if(w->right->color==BLACK&&w->left->color==BLACK){
					w->color=RED;
					x=x->parent;
				}else{
					if(w->left->color==BLACK){
						w->right->color=BLACK;
						w->color=RED;
						leftrotate(w);
						w=x->parent->left;
					}
					w->color=x->parent->color;
					x->parent->color=BLACK;
					w->left->color=BLACK;
					rightrotate(x->parent);
					x=root;
				}
			}
		}
		x->color=BLACK;
	}
	void updatesizup(node *x){
		while(x!=nil){
			updatesiz(x);
			x=x->parent;
		}
	}
	node* findnode(T key)const{
		node *x=root;
		while(x!=nil){
			if(key>x->key){
				x=x->right;
			}else if(key<x->key){
				x=x->left;
			}else{
				return x;
			}
		}
		return nil;
	}
public:
	RBTree(){
		nil=createnil();
		root=nil;
	}
	~RBTree(){
		clear();
		delete nil;
	}
	void clear(){
		clear(root);
		root=nil;
	}
	void clear(node *x){
		if(x==nil)return;
		clear(x->left);
		clear(x->right);
		delete x;
	}
	void insert(T key){
		node *y=nil;
		node *x=root;
		while(x!=nil){
			y=x;
			if(key<x->key){
				x=x->left;
			}else if(key>x->key){
				x=x->right;
			}else{
				x->cnt++;
				updatesizup(x);
				return;
			}
		}
		node *z=new node(key,RED,nil);
		z->parent=y;
		if(y==nil){
			root=z;
		}else if(key<y->key){
			y->left=z;
		}else{
			y->right=z;
		}
		updatesizup(z);
		insertfixup(z);
	}
	void remove(T key){
		node *z=findnode(key);
		if(z==nil)return;
		if(z->cnt>1){
			z->cnt--;
			updatesizup(z);
			return;
		}
		node *y=z;
		node *x;
		Color yoriginalcolor=y->color;
		if(z->left==nil){
			x=z->right;
			transplant(z,z->right);
			updatesizup(z->parent);
		}else if(z->right==nil){
			x=z->left;
			transplant(z,z->left);
			updatesizup(z->parent);
		}else{
			y=minimum(z->right);
			yoriginalcolor=y->color;
			x=y->right;
			if(y->parent==z){
				x->parent=y;
			}else{
				transplant(y,y->right);
				y->right=z->right;
				y->right->parent=y;
				updatesizup(y->parent);
			}
			transplant(z,y);
			y->left=z->left;
			y->left->parent=y;
			y->color=z->color;
			y->siz=z->siz;
			updatesiz(y);
			updatesizup(y);
		}
		if(yoriginalcolor==BLACK){
			deletefixup(x);
		}
		delete z;
	}
	T predecessor(T key)const{
		node *x=root;
		node *pred=nil;
		while(x!=nil){
			if(x->key<key){
				pred=x;
				x=x->right;
			}else{
				x=x->left;
			}
		}
		if(pred==nil){
			throw runtime_error("No predecessor exists");
		}
		return pred->key; 
	}
	T successor(T key)const{
		node *x=root;
		node *succ=nil;
		while(x!=nil){
			if(key<x->key){
				succ=x;
				x=x->left;
			}else{
				x=x->right;
			}
		}
		if(succ==nil){
			throw runtime_error("No successor exists");
		}
		return succ->key;
	}
	int rank(T key)const{
		int r=1;
		node *x=root;
		while(x!=nil){
			if(key<x->key){
				x=x->left;
			}else if(key>x->key){
				r+=x->left->siz+x->cnt; 
				x=x->right;
			}else{
				r+=x->left->siz;
				return r;
			}
		}
		return r;
	}
	T kth(int k)const{
		if(k<=0||k>root->siz){
			throw out_of_range("Rank out of range");
		}
		node *x=root;
		while(x!=nil){
			int leftsiz=x->left->siz;
			if(k<=leftsiz){
				x=x->left;
			}else if(k>leftsiz+x->cnt){
				k-=leftsiz+x->cnt;
				x=x->right;
			}else{
				return x->key;
			}
		}
		return T();
	}
	bool find(T key)const{
		return findnode(key)!=nil;
	}
	void printTree()const{
		if(root==nil){
			cout<<"Empty"<<endl;
			return;
		}
		std::queue<node*>q;
		q.push(root);
		while(!q.empty()){
			int lvlsiz=q.size();
			for(int i=0;i<lvlsiz;i++){
				node *x=q.front();
				q.pop();
				cout<<'('<<x->key<<(x->color==RED?'R':'B')<<"c:"<<x->cnt<<"s:"<<x->siz<<')'<<endl;
				if(x->left!=nil)q.push(x->left);
				if(x->right!=nil)q.push(x->right);
			}
			cout<<endl;
		}
	}
};
RBTree<int> rbt;
int n;
int main(){
	cin>>n;
	while(n--){
		int opt,x;
		cin>>opt>>x;
		switch(opt){
			case 1: rbt.insert(x);break;
			case 2: rbt.remove(x);break;
			case 3: cout<<rbt.rank(x)<<endl;break;
			case 4: cout<<rbt.kth(x)<<endl;break;
			case 5: cout<<rbt.predecessor(x)<<endl;break;
			case 6: cout<<rbt.successor(x)<<endl;break;
		}
	}
}

评测记录

为什么会死循环

2025/6/14 18:59
加载中...