--- tags: codebook --- {%hackmd theme-dark %} # AVL tree ```cpp= #include<iostream> #include<cstring> #include<queue> using namespace std; struct node{ node *ch[2]; int key, height; node(int x):ch({}),key(x),height(0){} void pull(); }; inline int h(node* cur){return cur?cur->height:-1;} void node::pull(){height=1+max(h(ch[0]),h(ch[1]));} node* rotate(node* cur,bool i){ node* tmp=cur->ch[i^1]; cur->ch[i^1]=tmp->ch[i^0]; tmp->ch[i^0]=cur,tmp->ch[i^1]->pull(); return tmp->pull(),tmp; } void balance(node*& cur){ for(auto i:{0,1}) if(h(cur->ch[i^1])-h(cur->ch[i^0])>1){ if(cur->ch[i^1]&&h(cur->ch[i^1]->ch[i^1])<h(cur->ch[i^1]->ch[i^0])) cur->ch[i^1]=rotate(cur->ch[i^1],i^1); cur=rotate(cur,i^0); } cur->pull(); } void insert(node*& cur,int x){ if(!cur) cur=new node(x); else insert(cur->ch[cur->key<x],x); balance(cur); } void print(node* cur){ if(!cur) return; queue<node*> q; q.push(cur); while(q.size()){ node* cur=q.front();q.pop(); cout<<cur->key<<' '; for(auto i:{0,1}) if(cur->ch[i]) q.push(cur->ch[i]); } cout<<endl; } void solve(){ int n; node* root=nullptr; while(cin>>n) insert(root,n); print(root); } int main(){ solve(); return 0; } ``` ```cpp= #include<iostream> #include<algorithm> #include<vector> using namespace std; template<typename T> struct avl_tree{ struct node{ T value; unsigned count;int height; node *left,*right; void update(){height=1+max(left->height,right->height);} node(const T& x=0):value(x),count(1),height(-1){} } *nul,*root; void init(node* cur){cur->left=cur->right=nul,cur->update();} void _print(ostream& os,node* cur){ if(cur==nul) return; _print(os,cur->left); os<<cur->value<<' '; _print(os,cur->right); } friend ostream& operator<<(ostream& os,avl_tree& t){ t._print(os,t.root); return os; } node* lrotate(node* cur){ node* tmp=cur->right; cur->right=tmp->left; tmp->left=cur; tmp->left->update(); tmp->update(); return tmp; } node* rrotate(node* cur){ node* tmp=cur->left; cur->left=tmp->right; tmp->right=cur; tmp->right->update(); tmp->update(); return tmp; } unsigned count(const T& x,node* cur){ if(cur==nul) return 0; else if(x==cur->value) return cur->count; else return count(x,(x<cur->value?cur->left:cur->right)); } unsigned count(const T& x){return count(x,root);} void update(node*& cur){ if(cur==nul) return; if(cur->right->height-cur->left->height>1){ if(cur->right->right->height<cur->right->left->height) cur->right=rrotate(cur->right); cur=lrotate(cur); }else if(cur->left->height-cur->right->height>1){ if(cur->left->left->height<cur->left->right->height) cur->left=lrotate(cur->left); cur=rrotate(cur); } cur->update(); } void insert(const T& x,node*& cur){ if(cur==nul) cur=new node(x),init(cur); else if(x==cur->value) cur->count++; else insert(x,(x<cur->value?cur->left:cur->right)); update(cur); } void insert(const T& x){insert(x,root);} void _erase(node*& cur,node*& pos){ if(pos==nul){ node* tmp=cur; cur=cur->left; delete tmp; }else if(pos->left==nul){ cur->value=pos->value; cur->count=pos->count; node* tmp=pos; pos=pos->right; delete tmp; }else _erase(cur,pos->left); update(pos); } bool erase_find(const T& x,node*& cur){ if(cur==nul) return false; else if(x==cur->value){_erase(cur,cur->right);return true;} else return erase_find(x,(x<cur->value?cur->left:cur->right)); update(cur); } bool erase(const T& x){return erase_find(x,root);} avl_tree(const vector<T>& v=vector<T>{}):nul(new node),root(nul){ for(const auto& i:v) insert(i); } }; template<typename T> ostream& operator<<(ostream& os,const vector<T>& v){ for(const auto& i:v) os<<i<<' '; return os; } int main(){ vector<int> vec{1,5,4,2,3,6,7,8,9,0}; avl_tree<int> tree(vec); cout<<vec<<endl; for(const auto& i:vec) cout<<tree<<endl,tree.erase(i); cout<<tree<<endl; return 0; } ```