template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::fix_up(Set::Node*root)const{if(is_red(root->rc)&&!is_red(root->lc))// fix right leaned red linkroot=rotate_left(root);if(is_red(root->lc)&&is_red(root->lc->lc))// fix doubly linked left leaned red link// if (root->lc == nullptr), then the second expr won't be evaluatedroot=rotate_right(root);if(is_red(root->lc)&&is_red(root->rc))// break up 4 nodecolor_flip(root);root->size=size(root->lc)+size(root->rc)+1;returnroot;}template<classKey,classCompare>typenameSet<Key,Compare>::Node_Set<Key,Compare>::insert(Set::Node_root,constKey&key)const{if(root==nullptr)returnnewNode(key,kRed,1);if(root->key==key);elseif(cmp\_(key,root->key))// if (key < root->key)root->lc=insert(root->lc,key);elseroot->rc=insert(root->rc,key);returnfix_up(root);}
template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::move_red_left(Set::Node*root)const{color_flip(root);if(is_red(root->rc->lc)){// assume that root->rc != nullptr when calling this functionroot->rc=rotate_right(root->rc);root=rotate_left(root);color_flip(root);}returnroot;}template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::delete_min(Set::Node*root)const{if(root->lc==nullptr){deleteroot;returnnullptr;}if(!is_red(root->lc)&&!is_red(root->lc->lc)){// make sure either root->lc or root->lc->lc is red// thus make sure we will delete a red node in the endroot=move_red_left(root);}root->lc=delete_min(root->lc);returnfix_up(root);}
删除任意节点
我们首先考虑删除叶子:与删最小值类似,我们在删除任意值的过程中也要维护一个性质,不过这次比较特殊,因为我们不是只向左边走,而是可以向左右两个方向走,因此在删除过程中维护的性质是这样的:如果往左走,当前节点是 h,那么需要保证 h 是红色,或者 h->lc 是红色;如果往右走,当前节点是 h,那么需要保证 h 是红色,或者 h->rc 是红色。这样可以保证我们最后总会删掉一个红色节点。
template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::delete_arbitrary(Set::Node*root,Keykey)const{if(cmp_(key,root->key)){// key < root->keyif(!is_red(root->lc)&&!(is_red(root->lc->lc)))root=move_red_left(root);// ensure the invariant: either root->lc or root->lc->lc (or root and// root->lc after dive into the function) is red, to ensure we will// eventually delete a red node. therefore we will not break the black// height balanceroot->lc=delete_arbitrary(root->lc,key);}else{// key >= root->keyif(is_red(root->lc))root=rotate_right(root);if(key==root->key&&root->rc==nullptr){deleteroot;returnnullptr;}if(!is_red(root->rc)&&!is_red(root->rc->lc))root=move_red_right(root);if(key==root->key){root->key=get_min(root->rc);root->rc=delete_min(root->rc);}else{root->rc=delete_arbitrary(root->rc,key);}}returnfix_up(root);}
#include<algorithm>#include<memory>#include<vector>template<classKey,classCompare=std::less<Key>>classSet{private:enumNodeColor{kBlack=0,kRed=1};structNode{Keykey;Node*lc{nullptr},*rc{nullptr};size_tsize{0};NodeColorcolor;// the color of the parent linkNode(Keykey,NodeColorcolor,size_tsize):key(key),color(color),size(size){}Node()=default;};voiddestroyTree(Node*root)const{if(root!=nullptr){destroyTree(root->lc);destroyTree(root->rc);root->lc=root->rc=nullptr;deleteroot;}}boolis_red(constNode*nd)const{returnnd==nullptr?false:nd->color;// kRed == 1, kBlack == 0}size_tsize(constNode*nd)const{returnnd==nullptr?0:nd->size;}Node*rotate_left(Node*node)const{// left rotate a red link// <1> <2>// / \\ // \ // * <2> ==> <1> *// / \ / \ // * * * *Node*res=node->rc;node->rc=res->lc;res->lc=node;res->color=node->color;node->color=kRed;res->size=node->size;node->size=size(node->lc)+size(node->rc)+1;returnres;}Node*rotate_right(Node*node)const{// right rotate a red link// <1> <2>// // \ / \\ // <2> * ==> * <1>// / \ / \ // * * * *Node*res=node->lc;node->lc=res->rc;res->rc=node;res->color=node->color;node->color=kRed;res->size=node->size;node->size=size(node->lc)+size(node->rc)+1;returnres;}NodeColorneg_color(NodeColorn)const{returnn==kBlack?kRed:kBlack;}voidcolor_flip(Node*node)const{node->color=neg_color(node->color);node->lc->color=neg_color(node->lc->color);node->rc->color=neg_color(node->rc->color);}Node*insert(Node*root,constKey&key)const;Node*delete_arbitrary(Node*root,Keykey)const;Node*delete_min(Node*root)const;Node*move_red_right(Node*root)const;Node*move_red_left(Node*root)const;Node*fix_up(Node*root)const;constKey&get_min(Node*root)const;voidserialize(Node*root,std::vector<Key>*)const;voidprint_tree(Set::Node*root,intindent)const;Comparecmp_=Compare();Node*root_{nullptr};public:usingKeyType=Key;usingValueType=Key;usingSizeType=std::size_t;usingDifferenceType=std::ptrdiff_t;usingKeyCompare=Compare;usingValueCompare=Compare;usingReference=Key&;usingConstReference=constKey&;Set()=default;Set(Set&)=default;Set(Set&&)noexcept=default;~Set(){destroyTree(root_);}SizeTypesize()const;SizeTypecount(constKeyType&key)const;SizeTypeerase(constKeyType&key);voidclear();voidinsert(constKeyType&key);boolempty()const;std::vector<Key>serialize()const;voidprint_tree()const;};template<classKey,classCompare>typenameSet<Key,Compare>::SizeTypeSet<Key,Compare>::count(ConstReferencekey)const{Node*x=root_;while(x!=nullptr){if(key==x->key)return1;if(cmp_(key,x->key))// if (key < x->key)x=x->lc;elsex=x->rc;}return0;}template<classKey,classCompare>typenameSet<Key,Compare>::SizeTypeSet<Key,Compare>::erase(constKeyType&key){if(count(key)>0){if(!is_red(root_->lc)&&!(is_red(root_->rc)))root_->color=kRed;root_=delete_arbitrary(root_,key);if(root_!=nullptr)root_->color=kBlack;return1;}else{return0;}}template<classKey,classCompare>voidSet<Key,Compare>::clear(){destroyTree(root_);root_=nullptr;}template<classKey,classCompare>voidSet<Key,Compare>::insert(constKeyType&key){root_=insert(root_,key);root_->color=kBlack;}template<classKey,classCompare>boolSet<Key,Compare>::empty()const{returnsize(root_)==0;}template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::insert(Set::Node*root,constKey&key)const{if(root==nullptr)returnnewNode(key,kRed,1);if(root->key==key);elseif(cmp_(key,root->key))// if (key < root->key)root->lc=insert(root->lc,key);elseroot->rc=insert(root->rc,key);returnfix_up(root);}template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::delete_min(Set::Node*root)const{if(root->lc==nullptr){deleteroot;returnnullptr;}if(!is_red(root->lc)&&!is_red(root->lc->lc)){// make sure either root->lc or root->lc->lc is red// thus make sure we will delete a red node in the endroot=move_red_left(root);}root->lc=delete_min(root->lc);returnfix_up(root);}template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::move_red_right(Set::Node*root)const{color_flip(root);if(is_red(root->lc->lc)){// assume that root->lc != nullptr when calling// this functionroot=rotate_right(root);color_flip(root);}returnroot;}template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::move_red_left(Set::Node*root)const{color_flip(root);if(is_red(root->rc->lc)){// assume that root->rc != nullptr when calling this functionroot->rc=rotate_right(root->rc);root=rotate_left(root);color_flip(root);}returnroot;}template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::fix_up(Set::Node*root)const{if(is_red(root->rc)&&!is_red(root->lc))// fix right leaned red linkroot=rotate_left(root);if(is_red(root->lc)&&is_red(root->lc->lc))// fix doubly linked left leaned red link// if (root->lc == nullptr), then the second expr won't be evaluatedroot=rotate_right(root);if(is_red(root->lc)&&is_red(root->rc))// break up 4 nodecolor_flip(root);root->size=size(root->lc)+size(root->rc)+1;returnroot;}template<classKey,classCompare>constKey&Set<Key,Compare>::get_min(Set::Node*root)const{Node*x=root;// will crash as intended when root == nullptrfor(;x->lc!=nullptr;x=x->lc);returnx->key;}template<classKey,classCompare>typenameSet<Key,Compare>::SizeTypeSet<Key,Compare>::size()const{returnsize(root_);}template<classKey,classCompare>typenameSet<Key,Compare>::Node*Set<Key,Compare>::delete_arbitrary(Set::Node*root,Keykey)const{if(cmp_(key,root->key)){// key < root->keyif(!is_red(root->lc)&&!(is_red(root->lc->lc)))root=move_red_left(root);// ensure the invariant: either root->lc or root->lc->lc (or root and// root->lc after dive into the function) is red, to ensure we will// eventually delete a red node. therefore we will not break the black// height balanceroot->lc=delete_arbitrary(root->lc,key);}else{// key >= root->keyif(is_red(root->lc))root=rotate_right(root);if(key==root->key&&root->rc==nullptr){deleteroot;returnnullptr;}if(!is_red(root->rc)&&!is_red(root->rc->lc))root=move_red_right(root);if(key==root->key){root->key=get_min(root->rc);root->rc=delete_min(root->rc);}else{root->rc=delete_arbitrary(root->rc,key);}}returnfix_up(root);}template<classKey,classCompare>std::vector<Key>Set<Key,Compare>::serialize()const{std::vector<int>v;serialize(root_,&v);returnv;}template<classKey,classCompare>voidSet<Key,Compare>::serialize(Set::Node*root,std::vector<Key>*res)const{if(root==nullptr)return;serialize(root->lc,res);res->push_back(root->key);serialize(root->rc,res);}template<classKey,classCompare>voidSet<Key,Compare>::print_tree(Set::Node*root,intindent)const{if(root==nullptr)return;print_tree(root->lc,indent+4);std::cout<<std::string(indent,'-')<<root->key<<std::endl;print_tree(root->rc,indent+4);}template<classKey,classCompare>voidSet<Key,Compare>::print_tree()const{print_tree(root_,0);}