AA Tree
概要
平衝二分探索木の一つ。
各ノードにlevelをもたせて管理し、各ノードは以下の制約を満たす。
-
葉ノードのlevelは 1
-
全てのノードのlevelはその左子ノードのlevelより 1 大きい
(node.level = node.left.level + 1
) -
全てのノードのlevelはその右子ノードのlevelと同じ、もしくは 1 大きい
(node.level = node.right.level
ornode.level = node.right.level + 1
) -
全てのノードのlevelはその右孫ノードのlevelより大きい
(node.level > node.right.right.level
) -
levelが 1 より大きい全てのノードは左右の子ノードを持つ
平衝処理は skew と split の2つの操作で実現される。
AA Tree は従来の赤黒木に「赤ノードは常に親の右子ノード」という制約を加えたものとみなせる。
計算量
各クエリ \(O(\log N)\)
実装
各ノードに以下の情報を持たせている
-
左右の子ノード, 親ノード
-
key, level
-
そのノード以下の部分木のサイズ
rotateに関する最適化はなし
#include<algorithm>
using namespace std;
template<typename T>
class AATree {
struct Node {
Node *left, *right, *prt;
T key;
int level, size;
Node(T x) {
left = right = prt = nullptr;
level = 1;
key = x;
size = 1;
}
inline Node* rotate_left() {
Node *r = this->right, *m = r->left, *p = this->prt;
if((r->prt = p)) {
if(p->left == this) p->left = r;
else p->right = r;
}
if((this->right = m)) m->prt = this;
r->left = this; this->prt = r;
int sz = this->size;
this->size += (m ? m->size : 0) - r->size;
r->size = sz;
return r;
}
inline Node* rotate_right() {
Node *l = this->left, *m = l->right, *p = this->prt;
if((l->prt = p)) {
if(p->left == this) p->left = l;
else p->right = l;
}
if((this->left = m)) m->prt = this;
l->right = this; this->prt = l;
int sz = this->size;
this->size += (m ? m->size : 0) - l->size;
l->size = sz;
return l;
}
inline bool is_left(Node *node) const {
return left == node;
}
inline void assign(Node *node) {
this->key = node->key;
}
};
Node *root;
inline Node* find_node(Node *node, T x) {
if(!node) return nullptr;
while(node->key != x) {
if(x < node->key) {
if(!node->left) break;
node = node->left;
} else {
if(!node->right) break;
node = node->right;
}
}
return node;
}
inline static Node* skew(Node *node) {
if(!node) return nullptr;
Node *left = node->left;
if(left && node->level == left->level) {
return node->rotate_right();
}
return node;
}
inline static Node* split(Node *node) {
if(!node) return nullptr;
Node *right = node->right;
if(right && right->right && node->level == right->right->level) {
Node* r = node->rotate_left();
r->level += 1;
return r;
}
return node;
}
inline static int get_level(Node *node) {
return (node ? node->level : 0);
}
public:
AATree() {
root = nullptr;
}
inline bool find(T x) {
Node *node = find_node(this->root, x);
return (node->key == x);
}
inline T at(int k) {
if(!this->root) {
return 0;
}
// assert(0 <= k < size);
Node *node = this->root;
++k;
while(1) {
int l_size = (node->left ? node->left->size : 0) + 1;
if(l_size == k) break;
if(k < l_size) {
node = node->left;
} else {
node = node->right;
k -= l_size;
}
}
return node->key;
}
inline bool insert(T x) {
if(!this->root) {
Node *new_node = new Node(x);
this->root = new_node;
return true;
}
Node *node = find_node(this->root, x);
if(node->key == x) return false;
Node *new_node = new Node(x);
if(x < node->key) {
node->left = new_node;
} else {
node->right = new_node;
}
new_node->prt = node;
node = new_node;
while(node) {
node = split(skew(node));
if(!node->prt) {
this->root = node;
break;
}
node = node->prt;
++node->size;
}
return true;
}
inline bool remove(T x) {
if(!this->root) return false;
Node *node = find_node(this->root, x);
if(node->key != x) return false;
if(node->left || node->right) {
Node *c_node = find_node(!node->left ? node->right : node->left, x);
node->assign(c_node);
node = c_node;
}
Node *prt = node->prt;
if(prt) {
if(prt->is_left(node)) {
prt->left = nullptr;
} else {
prt->right = nullptr;
}
} else {
this->root = nullptr;
}
node = prt;
while(node) {
// decrease_key(node)
int new_level = min(get_level(node->left), get_level(node->right)) + 1;
if(new_level < node->level) {
node->level = new_level;
if(new_level < get_level(node->right)) {
node->right->level = new_level;
}
}
--node->size;
node = skew(node);
if(skew(node->right)) skew(node->right->right);
node = split(node);
split(node->right);
if(!node->prt) {
this->root = node;
break;
}
node = node->prt;
}
return true;
}
int size() {
return (this->root ? this->root->size : 0);
}
};
Verified
-
AtCoder: "AtCoder Regular Contest 033 - C問題: データ構造": source (C++14, 131ms)
参考
-
Andersson, Arne. "Balanced search trees made simple." Workshop on Algorithms and Data Structures. Springer, Berlin, Heidelberg, 1993.