平衡树求调
查看原帖
平衡树求调
745907
szm_Only楼主2024/8/23 14:09

rt

#include<bits/stdc++.h>
using namespace std;
const int INF=0x7fffffff;

struct Node {
	int key;
	int h;
	Node* left;
	Node* right;
	int size;   
	int count;  
};

Node *_get(int key){
	Node *p=(Node *)malloc(sizeof(Node));
	p->key=key;
	p->h=1;
	p->left=p->right=nullptr;
	return p;
}

void _clr(Node *root){
	if(root==nullptr)return ;
	_clr(root->left);
	_clr(root->right);
	free(root);
	return ;
}

Node *pre(Node *root){
	Node *tmp=root->left;
	if(tmp==nullptr)return nullptr;
	while(!tmp->right)tmp=tmp->right;
	return tmp;
}

int _pre(Node *root,int x,int ans){
	if(root->key>=x){
		if(!root->left)return ans;
		else return _pre(root->left,x,ans);
	}else{
		if(!root->right)return root->key<x?root->key:ans;
		if(root->count!=0)return _pre(root->right,x,root->key);
		else return _pre(root->right,x,ans);
	}
}

int _nxt(Node *root,int x,int ans){
	if(root->key<=x){
		if(!root->right)return ans;
		else return _nxt(root->right,x,ans);
	}else{
		if(!root->left)return root->key<x?root->key:ans;
		if(root->count!=0)return _nxt(root->left,x,root->key);
		else return _nxt(root->left,x,ans);
	}
}

void _uph(Node *root){
	root->h=max(root->left?root->left->h+1:0,root->right?root->right->h+1:0);
	return ;
}

Node *l_rot(Node *root){
	Node *newroot=root->right;
	root->right=newroot->left;
	newroot->left=root;
	_uph(root);
	_uph(newroot);
	return newroot;
}

Node *r_rot(Node *root){
	Node *newroot=root->left;
	root->left=newroot->right;
	newroot->right=root;
	_uph(root);
	_uph(newroot);
	return newroot;	
}

Node *_mnt(Node *root){
	if(abs((root->left?root->left->h:0)-(root->right?root->right->h:0))<=1)return root;
	if(root->left->h>root->right->h){
		if(root->left->right->h>root->left->left->h){
			root->left=l_rot(root->left);
		}
		root=r_rot(root);
	}else{
		if(root->right->left->h>root->right->right->h){
			root->right=r_rot(root->right);
		}
		root=l_rot(root);		
	}
	return root;
}

Node *_ins(Node *root,int val){
	if(root==nullptr)return _get(val);
	if(val<root->key)_ins(root->left,val);
	else if(val>root->key)_ins(root->right,val);
	else root->count++;
	root->size=(root->count+
		(root->left?root->left->size:0)+
		(root->right?root->right->size:0)
		);
	_uph(root);
	return _mnt(root);
}

Node *_ers(Node *root,int k){
	if(root==nullptr)return root;
	if(k<root->key)root->left=_ers(root->left,k);
	else if(k>root->key)root->right=_ers(root->left,k);
	else {
		if(root->left==nullptr&&root->right==nullptr){
			free(root);
			return nullptr;
		}else if(root->left==nullptr||root->right==nullptr){
			Node *tmp=root->left?root->left:root->right;
			free(root);
			return tmp;
		}else{
			Node *tmp=pre(root);
			root->key=tmp->key;
			root->left=_ers(root->left,tmp->key);
		}
	}
	return _mnt(root);
}

int _rnk(Node *root,int val){
	if(root==nullptr)return 0;
	if(val==root->key)return (root->left?root->left->size:0)+1;
	else if(val<root->key)return _rnk(root->left,val);
	else return _rnk(root->right,val)+root->count+(root->left?root->left->size:0);
}

int _kth(Node *root,int k){
	if(root==nullptr)return -1;
	if(root->left){
		if(root->left->size>=k)return _kth(root->left,k);
		if(root->left->size>=k-root->count)return root->key;
	}else{
		if(k==1)return root->key;
	}
	return _kth(root->right,k-(root->left?root->left->size:0)-root->count);
}

int n,op,x;
Node *rt=nullptr;

int main(){
	scanf("%d",&n);
	while(n--){
		scanf("%d%d",&op,&x);
		if(op==1){
			rt=_ins(rt,x);
		}
		if(op==2){
			rt=_ers(rt,x);
		}
		if(op==3){
			printf("%d\n",_rnk(rt,x));
		}
		if(op==4){
			printf("%d\n",_kth(rt,x));
		}
		if(op==5){
			printf("%d\n",_pre(rt,x,-INF));
		}
		if(op==6){
			printf("%d\n",_nxt(rt,x,INF));
		}
		
	}
	return 0;
}
2024/8/23 14:09
加载中...