voidupdateSize(){USizeleftSize=this->left!=nullptr?this->left->size:0;USizerightSize=this->right!=nullptr?this->right->size:0;this->size=leftSize+rightSize+1;}staticvoidrotateLeft(NodePtr&node){assert(node!=nullptr);// clang-format off// | |// N S// / \ l-rotate(N) / \ // L S ==========> N R// / \ / \ // M R L M// clang-format onNodePtrsuccessor=node->right;node->right=successor->left;successor->left=node;node->updateSize();successor->updateSize();node=successor;}staticvoidrotateRight(NodePtr&node){assert(node!=nullptr);// clang-format off// | |// N S// / \ r-rotate(N) / \ // S R ==========> L N// / \ / \ // L M M R// clang-format onNodePtrsuccessor=node->left;node->left=successor->right;successor->right=node;node->updateSize();successor->updateSize();node=successor;}
维护
Case 1
size(N.left) < size(N.right.left)
1 2 3 4 5 6 7 8 910111213141516
if(size(node->right->left)>size(node->left)){// clang-format off// | | |// N N [M]// / \ r-rotate(R) / \ l-rotate(N) / \ // <L> R ==========> <L> [M] ==========> N R// / \ /// [M] R <L>// clang-format onrotateRight(node->right);rotateLeft(node);fixBalance(node->left);fixBalance(node->right);fixBalance(node);return;}
Case 2
size(N.left) < size(N.right.right)
1 2 3 4 5 6 7 8 91011121314
if(size(node->right->right)>size(node->left)){// clang-format off// | |// N R// / \ l-rotate(N) / \ // <L> R ==========> N [M]// \ /// [M] <L>// clang-format onrotateLeft(node);fixBalance(node->left);fixBalance(node);return;}
It can result in a destroyed SBT. But with the insertion above, a BST is still kept at the height of where is the total number of insertions, not the current size.
boolremove(NodePtr&node,Kkey,NodeConsumeraction){assert(node!=nullptr);if(key!=node->key){if(compare(key,node->key)){/* key < node->key */NodePtr&left=node->left;if(left!=nullptr&&remove(left,key,action)){node->updateSize();fixBalance(node);returntrue;}else{returnfalse;}}else{/* key > node->key */NodePtr&right=node->right;if(right!=nullptr&&remove(right,key,action)){node->updateSize();fixBalance(node);returntrue;}else{returnfalse;}}}assert(key==node->key);action(node);if(node->isLeaf()){// Case 1: no childnode=nullptr;}elseif(node->right==nullptr){// Case 2: left child only// clang-format off// P// | remove(N) P// N ========> |// / L// L// clang-format onnode=node->left;}elseif(node->left==nullptr){// Case 3: right child only// clang-format off// P// | remove(N) P// N ========> |// \ R// R// clang-format onnode=node->right;}elseif(node->right->left==nullptr){// Case 4: both left and right child, right child has no left child// clang-format off// | |// N remove(N) R// / \ ========> /// L R L// clang-format onNodePtrright=node->right;swapNode(node,right);right->right=node->right;node=right;node->updateSize();fixBalance(node);}else{// Case 5: both left and right child, right child is not a leaf// clang-format off// Step 1. find the node N with the smallest key// and its parent P on the right subtree// Step 2. swap S and N// Step 3. remove node N like Case 1 or Case 3// Step 4. update size for all nodes on the path// from S to P// | |// N S |// / \ / \ S// L .. swap(N, S) L .. remove(N) / \ // | =========> | ========> L ..// P P |// / \ / \ P// S .. N .. / \ // \ \ R ..// R R//// clang-format onstd::stack<NodePtr>path;// Step 1NodePtrsuccessor=node->right;NodePtrparent=node;path.push(node);while(successor->left!=nullptr){path.push(successor);parent=successor;successor=parent->left;}// Step 2swapNode(node,successor);// Step 3parent->left=node->right;// Restore nodenode=successor;// Step 4while(!path.empty()){path.top()->updateSize();path.pop();}}returntrue;}
/** * @file SizeBalancedTreeMap.hpp * @brief An SizeBalancedTree-based map implementation * @details The map is sorted according to the natural ordering of its * keys or by a {@code Compare} function provided; This implementation * provides guaranteed log(n) time cost for the contains, get, insert * and remove operations. * @author [r.ivance](https://github.com/RIvance) */#ifndef SIZE_BALANCED_TREE_MAP_HPP#define SIZE_BALANCED_TREE_MAP_HPP#include<cassert>#include<cstddef>#include<cstdint>#include<functional>#include<memory>#include<stack>#include<utility>#include<vector>/** * An SizeBalancedTree-based map implementation * http://wcipeg.com/wiki/Size_Balanced_Tree * @tparam Key the type of keys maintained by this map * @tparam Value the type of mapped values * @tparam Compare the compare function */template<typenameKey,typenameValue,typenameCompare=std::less<Key>>classSizeBalancedTreeMap{private:usingUSize=size_t;Comparecompare=Compare();public:structEntry{Keykey;Valuevalue;booloperator==(constEntry&rhs)constnoexcept{returnthis->key==rhs.key&&this->value==rhs.value;}booloperator!=(constEntry&rhs)constnoexcept{returnthis->key!=rhs.key||this->value!=rhs.value;}};private:structNode{usingPtr=std::shared_ptr<Node>;usingProvider=conststd::function<Ptr(void)>&;usingConsumer=conststd::function<void(constPtr&)>&;Keykey;Valuevalue{};Ptrleft=nullptr;Ptrright=nullptr;USizesize=1;explicitNode(Keyk):key(std::move(k)){}explicitNode(Keyk,Valuev):key(std::move(k)),value(std::move(v)){}~Node()=default;inlineboolisLeaf()constnoexcept{returnthis->left==nullptr&&this->right==nullptr;}inlinevoidupdateSize(){USizeleftSize=this->left!=nullptr?this->left->size:0;USizerightSize=this->right!=nullptr?this->right->size:0;this->size=leftSize+rightSize+1;}inlineEntryentry()const{returnEntry{key,value};}staticPtrfrom(constKey&k){returnstd::make_shared<Node>(Node(k));}staticPtrfrom(constKey&k,constValue&v){returnstd::make_shared<Node>(Node(k,v));}};usingNodePtr=typenameNode::Ptr;usingConstNodePtr=constNodePtr&;usingNodeProvider=typenameNode::Provider;usingNodeConsumer=typenameNode::Consumer;NodePtrroot=nullptr;usingK=constKey&;usingV=constValue&;public:usingEntryList=std::vector<Entry>;usingKeyValueConsumer=conststd::function<void(K,V)>&;usingMutKeyValueConsumer=conststd::function<void(K,Value&)>&;usingKeyValueFilter=conststd::function<bool(K,V)>&;classNoSuchMappingException:protectedstd::exception{private:constchar*message;public:explicitNoSuchMappingException(constchar*msg):message(msg){}constchar*what()constnoexceptoverride{returnmessage;}};SizeBalancedTreeMap()noexcept=default;/** * Returns the number of entries in this map. * @return size_t */inlineUSizesize()constnoexcept{if(this->root!=nullptr){returnthis->root->size;}else{return0;}}/** * Returns true if this collection contains no elements. * @return bool */inlineboolempty()constnoexcept{returnthis->root==nullptr;}/** * Removes all of the elements from this map. */voidclear()noexcept{this->root=nullptr;}/** * Returns the value to which the specified key is mapped; If this map * contains no mapping for the key, a {@code NoSuchMappingException} will * be thrown. * @param key * @return SizeBalancedTreeMap<Key, Value>::Value * @throws NoSuchMappingException */Valueget(Kkey)const{if(this->root==nullptr){throwNoSuchMappingException("Invalid key");}else{NodePtrnode=this->getNode(this->root,key);if(node!=nullptr){returnnode->value;}else{throwNoSuchMappingException("Invalid key");}}}/** * Returns the value to which the specified key is mapped; If this map * contains no mapping for the key, a new mapping with a default value * will be inserted. * @param key * @return SizeBalancedTreeMap<Key, Value>::Value & */Value&getOrDefault(Kkey){if(this->root==nullptr){this->root=Node::from(key);returnthis->root->value;}else{returnthis->getNodeOrProvide(this->root,key,[&key](){returnNode::from(key);})->value;}}/** * Returns true if this map contains a mapping for the specified key. * @param key * @return bool */boolcontains(Kkey)const{returnthis->getNode(this->root,key)!=nullptr;}/** * Associates the specified value with the specified key in this map. * @param key * @param value */voidinsert(Kkey,Vvalue){if(this->root==nullptr){this->root=Node::from(key,value);}else{this->insert(this->root,key,value);}}/** * If the specified key is not already associated with a value, associates * it with the given value and returns true, else returns false. * @param key * @param value * @return bool */boolinsertIfAbsent(Kkey,Vvalue){USizesizeBeforeInsertion=this->size();if(this->root==nullptr){this->root=Node::from(key,value);}else{this->insert(this->root,key,value,false);}returnthis->size()>sizeBeforeInsertion;}/** * If the specified key is not already associated with a value, associates * it with the given value and returns the value, else returns the associated * value. * @param key * @param value * @return SizeBalancedTreeMap<Key, Value>::Value & */Value&getOrInsert(Kkey,Vvalue){if(this->root==nullptr){this->root=Node::from(key,value);returnroot->value;}else{NodePtrnode=getNodeOrProvide(this->root,key,[&](){returnNode::from(key,value);});returnnode->value;}}Valueoperator[](Kkey)const{returnthis->get(key);}Value&operator[](Kkey){returnthis->getOrDefault(key);}/** * Removes the mapping for a key from this map if it is present; * Returns true if the mapping is present else returns false * @param key the key of the mapping * @return bool */boolremove(Kkey){if(this->root==nullptr){returnfalse;}else{returnthis->remove(this->root,key,[](ConstNodePtr){});}}/** * Removes the mapping for a key from this map if it is present and returns * the value which is mapped to the key; If this map contains no mapping for * the key, a {@code NoSuchMappingException} will be thrown. * @param key * @return SizeBalancedTreeMap<Key, Value>::Value * @throws NoSuchMappingException */ValuegetAndRemove(Kkey){Valueresult;NodeConsumeraction=[&](ConstNodePtrnode){result=node->value;};if(root==nullptr){throwNoSuchMappingException("Invalid key");}else{if(remove(this->root,key,action)){returnresult;}else{throwNoSuchMappingException("Invalid key");}}}/** * Gets the entry corresponding to the specified key; if no such entry * exists, returns the entry for the least key greater than the specified * key; if no such entry exists (i.e., the greatest key in the Tree is less * than the specified key), a {@code NoSuchMappingException} will be thrown. * @param key * @return SizeBalancedTreeMap<Key, Value>::Entry * @throws NoSuchMappingException */EntrygetCeilingEntry(Kkey)const{if(this->root==nullptr){throwNoSuchMappingException("No ceiling entry in this map");}NodePtrnode=this->root;std::stack<NodePtr>ancestors;while(node!=nullptr){if(key==node->key){returnnode->entry();}if(compare(key,node->key)){/* key < node->key */if(node->left!=nullptr){ancestors.push(node);node=node->left;}else{returnnode->entry();}}else{/* key > node->key */if(node->right!=nullptr){ancestors.push(node);node=node->right;}else{if(ancestors.empty()){throwNoSuchMappingException("No ceiling entry in this map");}NodePtrparent=ancestors.top();ancestors.pop();while(node==parent->right){node=parent;if(!ancestors.empty()){parent=ancestors.top();ancestors.pop();}else{throwNoSuchMappingException("No ceiling entry in this map");}}returnparent->entry();}}}throwNoSuchMappingException("No ceiling entry in this map");}/** * Gets the entry corresponding to the specified key; if no such entry exists, * returns the entry for the greatest key less than the specified key; * if no such entry exists, a {@code NoSuchMappingException} will be thrown. * @param key * @return SizeBalancedTreeMap<Key, Value>::Entry * @throws NoSuchMappingException */EntrygetFloorEntry(Kkey)const{if(this->root==nullptr){throwNoSuchMappingException("No floor entry exists in this map");}NodePtrnode=this->root;std::stack<NodePtr>ancestors;while(node!=nullptr){if(key==node->key){returnnode->entry();}if(compare(key,node->key)){/* key < node->key */if(node->left!=nullptr){ancestors.push(node);node=node->left;}else{if(ancestors.empty()){throwNoSuchMappingException("No floor entry exists in this map");}NodePtrparent=ancestors.top();ancestors.pop();while(node==parent->left){node=parent;if(!ancestors.empty()){parent=ancestors.top();ancestors.pop();}else{throwNoSuchMappingException("No floor entry exists in this map");}}returnparent->entry();}}else{/* key > node->key */if(node->right!=nullptr){ancestors.push(node);node=node->right;}else{returnnode->entry();}}}throwNoSuchMappingException("No floor entry exists in this map");}/** * Gets the entry for the least key greater than the specified * key; if no such entry exists, returns the entry for the least * key greater than the specified key; if no such entry exists, * a {@code NoSuchMappingException} will be thrown. * @param key * @return SizeBalancedTreeMap<Key, Value>::Entry * @throws NoSuchMappingException */EntrygetHigherEntry(Kkey){if(this->root==nullptr){throwNoSuchMappingException("No higher entry exists in this map");}NodePtrnode=this->root;std::stack<NodePtr>ancestors;while(node!=nullptr){if(compare(key,node->key)){/* key < node->key */if(node->left!=nullptr){ancestors.push(node);node=node->left;}else{returnnode->entry();}}else{/* key >= node->key */if(node->right!=nullptr){ancestors.push(node);node=node->right;}else{if(ancestors.empty()){throwNoSuchMappingException("No higher entry exists in this map");}NodePtrparent=ancestors.top();ancestors.pop();while(node==parent->right){node=parent;if(!ancestors.empty()){parent=ancestors.top();ancestors.pop();}else{throwNoSuchMappingException("No higher entry exists in this map");}}returnparent->entry();}}}throwNoSuchMappingException("No higher entry exists in this map");}/** * Returns the entry for the greatest key less than the specified key; if * no such entry exists (i.e., the least key in the Tree is greater than * the specified key), a {@code NoSuchMappingException} will be thrown. * @param key * @return SizeBalancedTreeMap<Key, Value>::Entry * @throws NoSuchMappingException */EntrygetLowerEntry(Kkey)const{if(this->root==nullptr){throwNoSuchMappingException("No lower entry exists in this map");}NodePtrnode=this->root;std::stack<NodePtr>ancestors;while(node!=nullptr){if(compare(key,node->key)||key==node->key){/* key <= node->key */if(node->left!=nullptr){ancestors.push(node);node=node->left;}else{if(ancestors.empty()){throwNoSuchMappingException("No lower entry exists in this map");}NodePtrparent=ancestors.top();ancestors.pop();while(node==parent->left){node=parent;if(!ancestors.empty()){parent=ancestors.top();ancestors.pop();}else{throwNoSuchMappingException("No lower entry exists in this map");}}returnparent->entry();}}else{/* key > node->key */if(node->right!=nullptr){ancestors.push(node);node=node->right;}else{returnnode->entry();}}}throwNoSuchMappingException("No lower entry exists in this map");}/** * Count the number of entries that are less than the given key * @param key * @return USize */USizecountLessThan(Kkey){returnthis->countLess(this->root,key);}/** * Count the number of entries that are less or equal to the given key * @param key * @return USize */USizecountLessOrEqualTo(Kkey){returnthis->countLess(this->root,key,true);}/** * Count the number of entries that are greater than the given key * @param key * @return USize */USizecountGreaterThan(Kkey){returnthis->countGreater(this->root,key);}/** * Count the number of entries that are greater or equal to the given key * @param key * @return USize */USizecountGreaterOrEqualTo(Kkey){returnthis->countGreater(this->root,key,true);}/** * Remove all entries that satisfy the filter condition. * @param filter */voidremoveAll(KeyValueFilterfilter){std::vector<Key>keys;this->inorderTraversal([&](ConstNodePtrnode){if(filter(node->key,node->value)){keys.push_back(node->key);}});for(constKey&key:keys){this->remove(key);}}/** * Performs the given action for each key and value entry in this map. * The value is immutable for the action. * @param action */voidforEach(KeyValueConsumeraction)const{this->inorderTraversal([&](ConstNodePtrnode){action(node->key,node->value);});}/** * Performs the given action for each key and value entry in this map. * The value is mutable for the action. * @param action */voidforEachMut(MutKeyValueConsumeraction){this->inorderTraversal([&](ConstNodePtrnode){action(node->key,node->value);});}/** * Returns a list containing all of the entries in this map. * @return SizeBalancedTreeMap<Key, Value>::EntryList */EntryListtoEntryList()const{EntryListentryList;this->inorderTraversal([&](ConstNodePtrnode){entryList.push_back(node->entry());});returnentryList;}private:staticUSizesize(ConstNodePtrnode){returnnode!=nullptr?node->size:0;}staticvoidrotateLeft(NodePtr&node){assert(node!=nullptr);// clang-format off// | |// N S// / \ l-rotate(N) / \ // L S ==========> N R// / \ / \ // M R L M// clang-format onNodePtrsuccessor=node->right;node->right=successor->left;successor->left=node;node->updateSize();successor->updateSize();node=successor;}staticvoidrotateRight(NodePtr&node){assert(node!=nullptr);// clang-format off// | |// N S// / \ r-rotate(N) / \ // S R ==========> L N// / \ / \ // L M M R// clang-format onNodePtrsuccessor=node->left;node->left=successor->right;successor->right=node;node->updateSize();successor->updateSize();node=successor;}staticvoidswapNode(NodePtr&lhs,NodePtr&rhs){std::swap(lhs->key,rhs->key);std::swap(lhs->value,rhs->value);std::swap(lhs,rhs);}staticvoidfixBalance(NodePtr&node){if(node==nullptr){return;}if(node->left!=nullptr){if(size(node->left->left)>size(node->right)){// clang-format off// | |// N L// / \ r-rotate(N) / \ // L <R> ==========> [M] N// / \ // [M] <R>// clang-format onrotateRight(node);fixBalance(node->right);fixBalance(node);return;}elseif(size(node->left->right)>size(node->right)){// clang-format off// | | |// N N [M]// / \ l-rotate(L) / \ r-rotate(N) / \ // L <R> ==========> [M] <R> ==========> L N// \ / \ // [M] L <R>// clang-format onrotateLeft(node->left);rotateRight(node);fixBalance(node->left);fixBalance(node->right);fixBalance(node);return;}}if(node->right!=nullptr){if(size(node->right->right)>size(node->left)){// clang-format off// | |// N R// / \ l-rotate(N) / \ // <L> R ==========> N [M]// \ /// [M] <L>// clang-format onrotateLeft(node);fixBalance(node->left);fixBalance(node);return;}elseif(size(node->right->left)>size(node->left)){// clang-format off// | | |// N N [M]// / \ r-rotate(R) / \ l-rotate(N) / \ // <L> R ==========> <L> [M] ==========> N R// / \ /// [M] R <L>// clang-format onrotateRight(node->right);rotateLeft(node);fixBalance(node->left);fixBalance(node->right);fixBalance(node);return;}}}NodePtrgetNodeOrProvide(NodePtr&node,Kkey,NodeProviderprovide){assert(node!=nullptr);if(key==node->key){returnnode;}assert(key!=node->key);NodePtrresult;if(compare(key,node->key)){/* key < node->key */if(node->left==nullptr){result=node->left=provide();node->updateSize();}else{result=getNodeOrProvide(node->left,key,provide);node->updateSize();fixBalance(node);}}else{/* key > node->key */if(node->right==nullptr){result=node->right=provide();node->updateSize();}else{result=getNodeOrProvide(node->right,key,provide);node->updateSize();fixBalance(node);}}returnresult;}NodePtrgetNode(ConstNodePtrnode,Kkey)const{assert(node!=nullptr);if(key==node->key){returnnode;}if(compare(key,node->key)){/* key < node->key */returnnode->left==nullptr?nullptr:getNode(node->left,key);}else{/* key > node->key */returnnode->right==nullptr?nullptr:getNode(node->right,key);}}voidinsert(NodePtr&node,Kkey,Vvalue,boolreplace=true){assert(node!=nullptr);if(key==node->key){if(replace){node->value=value;}return;}assert(key!=node->key);if(compare(key,node->key)){/* key < node->key */if(node->left==nullptr){node->left=Node::from(key,value);node->updateSize();}else{insert(node->left,key,value,replace);node->updateSize();fixBalance(node);}}else{/* key > node->key */if(node->right==nullptr){node->right=Node::from(key,value);node->updateSize();}else{insert(node->right,key,value,replace);node->updateSize();fixBalance(node);}}}boolremove(NodePtr&node,Kkey,NodeConsumeraction){assert(node!=nullptr);if(key!=node->key){if(compare(key,node->key)){/* key < node->key */NodePtr&left=node->left;if(left!=nullptr&&remove(left,key,action)){node->updateSize();fixBalance(node);returntrue;}else{returnfalse;}}else{/* key > node->key */NodePtr&right=node->right;if(right!=nullptr&&remove(right,key,action)){node->updateSize();fixBalance(node);returntrue;}else{returnfalse;}}}assert(key==node->key);action(node);if(node->isLeaf()){// Case 1: no childnode=nullptr;}elseif(node->right==nullptr){// clang-format off// Case 2: left child only// P// | remove(N) P// N ========> |// / L// L// clang-format onnode=node->left;}elseif(node->left==nullptr){// clang-format off// Case 3: right child only// P// | remove(N) P// N ========> |// \ R// R// clang-format onnode=node->right;}elseif(node->right->left==nullptr){// clang-format off// Case 4: both left and right child, right child has no left child// | |// N remove(N) R// / \ ========> /// L R L// clang-format onNodePtrright=node->right;swapNode(node,right);right->right=node->right;node=right;node->updateSize();fixBalance(node);}else{// clang-format off// Case 5: both left and right child, right child is not a leaf// Step 1. find the node N with the smallest key// and its parent P on the right subtree// Step 2. swap S and N// Step 3. remove node N like Case 1 or Case 3// Step 4. update size for all nodes on the path// from S to P// | |// N S |// / \ / \ S// L .. swap(N, S) L .. remove(N) / \ // | =========> | ========> L ..// P P |// / \ / \ P// S .. N .. / \ // \ \ R ..// R R// clang-format onstd::stack<NodePtr>path;// Step 1NodePtrsuccessor=node->right;NodePtrparent=node;path.push(node);while(successor->left!=nullptr){path.push(successor);parent=successor;successor=parent->left;}// Step 2swapNode(node,successor);// Step 3parent->left=node->right;// Restore nodenode=successor;// Step 4while(!path.empty()){path.top()->updateSize();path.pop();}}returntrue;}USizecountLess(ConstNodePtrnode,Kkey,boolcountEqual=false)const{if(node==nullptr){return0;}elseif(key<node->key){returncountLess(node->left,key,countEqual);}elseif(key>node->key){returnsize(node->left)+1+countLess(node->right,key,countEqual);}else{returnsize(node->left)+(countEqual?1:0);}}USizecountGreater(ConstNodePtrnode,Kkey,boolcountEqual=false)const{if(node==nullptr){return0;}elseif(key<node->key){returnsize(node->right)+1+countGreater(node->left,key,countEqual);}elseif(key>node->key){returncountGreater(node->right,key,countEqual);}else{returnsize(node->right)+(countEqual?1:0);}}voidinorderTraversal(NodeConsumeraction)const{if(this->root==nullptr){return;}std::stack<NodePtr>stack;NodePtrnode=this->root;while(node!=nullptr||!stack.empty()){while(node!=nullptr){stack.push(node);node=node->left;}if(!stack.empty()){node=stack.top();stack.pop();action(node);node=node->right;}}}};#endif // SIZE_BALANCED_TREE_MAP_HPP