AVL thụ

/**
* @brief An AVLTree-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.
*/

#ifndef AVLTREE_MAP_HPP
#define AVLTREE_MAP_HPP

#include&LTcassert>
#include&LTcstddef>
#include&LTcstdint>
#include&LTfunctional>
#include&LTmemory>
#include&LTstack>
#include&LTutility>
#include&LTvector>

/**
* An AVLTree-based map implementation
* https://en.wikipedia.org/wiki/AVL_tree
* @tparam Key the type of keys maintained by this map
* @tparam Value the type of mapped values
* @tparam Compare
*/
templatetypenameKey,typenameValue,typenameCompare=std::lessKey>>
classAvlTreeMap{
private:
usingUSize=size_t;
usingFactor=int64_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_ptrNode>;
usingProvider=conststd::functionPtr(void)>&;
usingConsumer=conststd::functionvoid(constPtr&)>&;

Keykey;
Valuevalue{};

Ptrleft=nullptr;
Ptrright=nullptr;

USizeheight=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;
}

inlinevoidupdateHeight()noexcept{
if(this->isLeaf()){
this->height=1;
}elseif(this->left==nullptr){
this->height=this->right->height+1;
}elseif(this->right==nullptr){
this->height=this->left->height+1;
}else{
this->height=std::max(left->height,right->height)+1;
}
}

inlineFactorfactor()constnoexcept{
if(this->isLeaf()){
return0;
}elseif(this->left==nullptr){
return(Factor)this->right->height;
}elseif(this->right==nullptr){
return(Factor)-this->left->height;
}else{
return(Factor)(this->right->height-this->left->height);
}
}

inlineEntryentry()const{returnEntry{key,value};}

staticPtrfrom(constKey&k){returnstd::make_sharedNode>(Node(k));}

staticPtrfrom(constKey&k,constValue&v){
returnstd::make_sharedNode>(Node(k,v));
}
};

usingNodePtr=typenameNode::Ptr;
usingConstNodePtr=constNodePtr&;
usingNodeProvider=typenameNode::Provider;
usingNodeConsumer=typenameNode::Consumer;

NodePtrroot=nullptr;
USizecount=0;

usingK=constKey&;
usingV=constValue&;

public:
usingEntryList=std::vectorEntry>;
usingKeyValueConsumer=conststd::functionvoid(K,V)>&;
usingMutKeyValueConsumer=conststd::functionvoid(K,Value&)>&;
usingKeyValueFilter=conststd::functionbool(K,V)>&;

classNoSuchMappingException:protectedstd::exception{
private:
constchar*message;

public:
explicitNoSuchMappingException(constchar*msg):message(msg){}

constchar*what()constnoexceptoverride{returnmessage;}
};

AvlTreeMap()noexcept=default;

/**
* Returns the number of entries in this map.
* @return size_t
*/
inlineUSizesize()constnoexcept{returnthis->count;}

/**
* Returns true if this collection contains no elements.
* @return bool
*/
inlineboolempty()constnoexcept{returnthis->count==0;}

/**
* Removes all of the elements from this map.
*/
voidclear()noexcept{
this->root=nullptr;
this->count=0;
}

/**
* 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 AvlTreeMap&LTKey, 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 AvlTreeMap&LTKey, Value>::Value &
*/
Value&getOrDefault(Kkey){
if(this->root==nullptr){
this->root=Node::from(key);
this->count+=1;
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);
this->count+=1;
}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);
this->count+=1;
}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
*/
Value&getOrInsert(Kkey,Vvalue){
if(this->root==nullptr){
this->root=Node::from(key,value);
this->count+=1;
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 AvlTreeMap&LTKey, 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 AvlTreeMap&LTKey, Value>::Entry
* @throws NoSuchMappingException
*/
EntrygetCeilingEntry(Kkey)const{
if(this->root==nullptr){
throwNoSuchMappingException("No ceiling entry in this map");
}

NodePtrnode=this->root;
std::stackNodePtr>ancestors;

while(node!=nullptr){
if(key==node->key){
returnnode->entry();
}

if(compare(key,node->key)){
/* key 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 AvlTreeMap&LTKey, Value>::Entry
* @throws NoSuchMappingException
*/
EntrygetFloorEntry(Kkey)const{
if(this->root==nullptr){
throwNoSuchMappingException("No floor entry exists in this map");
}

NodePtrnode=this->root;
std::stackNodePtr>ancestors;

while(node!=nullptr){
if(key==node->key){
returnnode->entry();
}

if(compare(key,node->key)){
/* key 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 AvlTreeMap&LTKey, Value>::Entry
* @throws NoSuchMappingException
*/
EntrygetHigherEntry(Kkey){
if(this->root==nullptr){
throwNoSuchMappingException("No higher entry exists in this map");
}

NodePtrnode=this->root;
std::stackNodePtr>ancestors;

while(node!=nullptr){
if(compare(key,node->key)){
/* key 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 AvlTreeMap&LTKey, Value>::Entry
* @throws NoSuchMappingException
*/
EntrygetLowerEntry(Kkey)const{
if(this->root==nullptr){
throwNoSuchMappingException("No lower entry exists in this map");
}

NodePtrnode=this->root;
std::stackNodePtr>ancestors;

while(node!=nullptr){
if(compare(key,node->key)||key==node->key){
/* key 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");
}

/**
* Remove all entries that satisfy the filter condition.
* @param filter
*/
voidremoveAll(KeyValueFilterfilter){
std::vectorKey>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 AvlTreeMap&LTKey, Value>::EntryList
*/
EntryListtoEntryList()const{
EntryListentryList;
this->inorderTraversal(
[&](ConstNodePtrnode){entryList.push_back(node->entry());});
returnentryList;
}

private:
staticNodePtrrotateLeft(ConstNodePtrnode){
// clang-format off
// | |
// N S
// / \ l-rotate(N) / \
// L S ==========> N R
// / \ / \
// M R L M
NodePtrsuccessor=node->right;
// clang-format on
node->right=successor->left;
successor->left=node;

node->updateHeight();
successor->updateHeight();

returnsuccessor;
}

staticNodePtrrotateRight(ConstNodePtrnode){
// clang-format off
// | |
// N S
// / \ r-rotate(N) / \
// S R ==========> L N
// / \ / \
// L M M R
NodePtrsuccessor=node->left;
// clang-format on
node->left=successor->right;
successor->right=node;

node->updateHeight();
successor->updateHeight();

returnsuccessor;
}

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->factor()-1){
if(node->left->factor()0){
// clang-format off
// Left-Left Case
// |
// C |
// / r-rotate(C) B
// B ==========> / \
// / A C
// A
// clang-format on
node=rotateRight(node);
}else{
// clang-format off
// Left-Right Case
// | |
// C C |
// / l-rotate(A) / r-rotate(C) B
// A ==========> B ==========> / \
// \ / A C
// B A
// clang-format on
node->left=rotateLeft(node->left);
node=rotateRight(node);
}
}elseif(node->factor()>1){
if(node->right->factor()>0){
// clang-format off
// Right-Right Case
// |
// C |
// \ l-rotate(C) B
// B ==========> / \
// \ A C
// A
// clang-format on
node=rotateLeft(node);
}else{
// clang-format off
// Right-Left Case
// | |
// A A |
// \ r-rotate(C) \ l-rotate(A) B
// C ==========> B ==========> / \
// / \ A C
// B C
// clang-format on
node->right=rotateRight(node->right);
node=rotateLeft(node);
}
}
}

NodePtrgetNodeOrProvide(NodePtr&node,Kkey,NodeProviderprovide){
assert(node!=nullptr);

if(key==node->key){
returnnode;
}

assert(key!=node->key);

NodePtrresult;

if(compare(key,node->key)){
/* key key */
if(node->left==nullptr){
result=node->left=provide();
this->count+=1;
node->updateHeight();
}else{
result=getNodeOrProvide(node->left,key,provide);
node->updateHeight();
fixBalance(node);
}
}else{
/* key > node->key */
if(node->right==nullptr){
result=node->right=provide();
this->count+=1;
node->updateHeight();
}else{
result=getNodeOrProvide(node->right,key,provide);
node->updateHeight();
fixBalance(node);
}
}

returnresult;
}

NodePtrgetNode(ConstNodePtrnode,Kkey)const{
assert(node!=nullptr);

if(key==node->key){
returnnode;
}

if(compare(key,node->key)){
/* key 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 key */
if(node->left==nullptr){
node->left=Node::from(key,value);
this->count+=1;
node->updateHeight();
}else{
insert(node->left,key,value,replace);
node->updateHeight();
fixBalance(node);
}
}else{
/* key > node->key */
if(node->right==nullptr){
node->right=Node::from(key,value);
this->count+=1;
node->updateHeight();
}else{
insert(node->right,key,value,replace);
node->updateHeight();
fixBalance(node);
}
}
}

boolremove(NodePtr&node,Kkey,NodeConsumeraction){
assert(node!=nullptr);

if(key!=node->key){
if(compare(key,node->key)){
/* key key */
NodePtr&left=node->left;
if(left!=nullptr&&remove(left,key,action)){
node->updateHeight();
fixBalance(node);
returntrue;
}else{
returnfalse;
}
}else{
/* key > node->key */
NodePtr&right=node->right;
if(right!=nullptr&&remove(right,key,action)){
node->updateHeight();
fixBalance(node);
returntrue;
}else{
returnfalse;
}
}
}

assert(key==node->key);
action(node);

if(node->isLeaf()){
// Case 1: no child
node=nullptr;
}elseif(node->right==nullptr){
// clang-format off
// Case 2: left child only
// P
// | remove(N) P
// N ========> |
// / L
// L
// clang-format on
node=node->left;
node->updateHeight();
}elseif(node->left==nullptr){
// clang-format off
// Case 3: right child only
// P
// | remove(N) P
// N ========> |
// \ R
// R
// clang-format on
node=node->right;
node->updateHeight();
}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 on
NodePtrright=node->right;
swapNode(node,right);
right->right=node->right;
node=right;
node->updateHeight();
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 height for P
// | |
// N S |
// / \ / \ S
// L.. swap(N, S) L.. remove(N) / \
// | =========> | ========> L..
// P P |
// / \ / \ P
// S.. N.. / \
// \ \ R..
// R R
// clang-format on

// Step 1
NodePtrsuccessor=node->right;
NodePtrparent=node;
while(successor->left!=nullptr){
parent=successor;
successor=parent->left;
}
// Step 2
swapNode(node,successor);
// Step 3
parent->left=node->right;
// Restore node
node=successor;
// Step 4
parent->updateHeight();
}

this->count-=1;
returntrue;
}

voidinorderTraversal(NodeConsumeraction)const{
if(this->root==nullptr){
return;
}

std::stackNodePtr>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// AVLTREE_MAP_HPP