この記事はCompetitive Programming (1) Advent Calendar 7日目の記事とISer Advent Calendar 8日目の記事として書かれました。
TL;DR
Implicit Treapを実装しました。以下の操作がクエリ毎O(logn)で可能です。
- 配列のランダムアクセス
- 配列の任意箇所へのinsert
- 任意要素のerase
- 任意区間に対するsum, max, min等のクエリ
- 任意区間に対する加算、更新等のクエリ
- 任意区間のreverse
- 任意区間のrotate
- Implicit Treap上の二分探索
この記事では一般的なTreapから順を追って説明していこうと思うのでImplicit Treapだけコピペして使いたい!という人は最後の方へどうぞ。
追記
bug fixとかをした最新の実装はここにあります。
- [12/24] [] operatorと累積を返すクエリで長さが0になるときのバグを修正しました。
前提知識
二分探索木に関する性質は既知とします。Segment TreeやBinary Indexed Treeに関する知識(遅延評価を含める)があれば理解にはかなり役立ちますが、まあなくても読めるかなと思います。Treapは知らなくても大丈夫です。
Treap (Cartesian Tree)とは
Treapという名前はtree + heapの略で、ペア(X, Y)を格納するデータ構造なんですが、Xに関しては二分探索木、Yに関しては最大値ヒープになっているもののことをいいます。Xはkey, Yはpriorityと呼ばれます。愚直な二分探索木は挿入等の操作に最悪O(n)を要しますが、ランダムなpriorityを導入することにより確率的にO(logn)の計算量を達成しようというものです(どう活用されているかは下で説明されます)。
Treapの操作と実装方法
search(X)
通常の二分探索木のsearchと同様です。O(logn)
split(T, X)
Treapを分割して2つのTreapを作ります。片方はX未満のkeyからなるもので、もう一方はX以上のkeyからなるものです。これは部分木に関する再帰でO(Tの深さ)で書くことができます。
void split(Tree t, T key, Tree& l, Tree& r) { if (!t) { l = r = nullptr; } else if (key < t->key) { split(t->l, key, l, t->l), r = t; } else { split(t->r, key, t->r, r), l = t; } }
insert(X, Y)
通常の二分探索木の挿入の操作と同様に、key Xに注目して木を降りていくのですが、priorityがY未満になっているノードVに遭遇したらストップします。将来的にこのノードを挿入したいノードに置き換えるのですが、それにはsplit操作を用います。すなわち、Vの部分木を値Xでsplitして得られる2つの部分木をL, Rとして、Vを(X, Y)に置き換えたのち、(X, Y)の左の子の部分木をL, 右の子の部分木をRにすればよいです。計算量はO(logn)です。
void insert(Tree& t, Tree item) { if (!t) { t = item; } else if (item->priority > t->priority) { split(t, item->key, item->l, item->r), t = item; } else { insert(item->key < t->key ? t->l : t->r, item); } }
merge(T1, T2)
2つのTreapをマージする操作です。これには、priorityがより高いrootを選び、その部分木に残りを再帰的にマージしていく、という再帰的な操作で書くことができます。オーダーはO(Tの深さ)です。
void merge(Tree& t, Tree l, Tree r) { if (!l || !r) { t = l ? l : r; } else if (l->priority > r->priority) { merge(l->r, l->r, r), t = l; } else { merge(r->l, l, r->l), t = r; } }
erase(X)
通常の二分探索木のsearchと同様にXのkeyをもつノードVに行きます。Vの部分木全体を、Vの左と子の部分木と右の子の部分木をmergeした木で置き換えれば終了です。計算量はO(logn)です。
void erase(Tree& t, T key) { if (t->key == key) { merge(t, t->l, t->r); } else { erase(key < t->key ? t->l : t->r, key); } }
Treapの実装
template <class T> class Treap { xorshift rnd; struct Node { T key; int priority; Node *l, *r; Node(T key, int priority) : key(key), priority(priority), l(nullptr), r(nullptr) {} } *root = nullptr; using Tree = Node *; void split(Tree t, T key, Tree& l, Tree& r) { if (!t) { l = r = nullptr; } else if (key < t->key) { split(t->l, key, l, t->l), r = t; } else { split(t->r, key, t->r, r), l = t; } } void merge(Tree& t, Tree l, Tree r) { if (!l || !r) { t = l ? l : r; } else if (l->priority > r->priority) { merge(l->r, l->r, r), t = l; } else { merge(r->l, l, r->l), t = r; } } void insert(Tree& t, Tree item) { if (!t) { t = item; } else if (item->priority > t->priority) { split(t, item->key, item->l, item->r), t = item; } else { insert(item->key < t->key ? t->l : t->r, item); } } void erase(Tree& t, T key) { if (t->key == key) { merge(t, t->l, t->r); } else { erase(key < t->key ? t->l : t->r, key); } } bool find(Tree& t, T key) { if (!t) { return false; } else if (t->key == key) { return true; } else { return find(key < t->key ? t->l : t->r, key); } } public: void insert(T key) { insert(root, new Node(key, rnd.random())); } void erase(T key) { erase(root, key); } bool find(T key) { return find(root, key); } };
こんな感じにわりとシンプルに書けます。Treapの概要がわかったところでいよいよ本題のImplicit Treapに進みます。
Implicit Treap Revisited
Implicit Treapは、ユーザには配列のようなインターフェースを提供しますが、中身はTreapです。インデックスがkの要素に対応するNodeはkというimplicit keyをもち、Node間の大小関係はそれらのimplicit keyの大小関係として定義されます。
実は、本記事で実装されるImplicit Treapよりも機能が制限されたバージョンがC++のSTLにも存在し、rope
と呼ばれています(ext/rope
をインクルードすると使えます)。
Implicit cartesian tree in GNU C++ STL. - Codeforces
Implicit Treapにおける区間の扱い
Segment TreeやBinary Indexed Treeでは、各Nodeにその部分木についてのmin/sum等の情報を保存しました。というか各Nodeにある区間が対応付けられていました。区間に対するクエリに答える際は、その区間をちょうど被覆するようなNodeの集合を見つけ、それらのmin/sum等について計算を行うことでクエリを処理しました。
Implicit Treapでも同様の考え方をします。なんならImplicit Treapのほうがシンプルです。Implicit Treapでは区間は、あるノードの部分木として表されます。Split操作により、木を分けることが可能なので、ある区間を(Segment Treeのときのようにいくつかの区間の被覆ではなく)一つの区間としてまるごと取り出せます。すなわち、[0, N)から[L, R)を切り出したいときは、まず[0, R), [R, N)にSplitし、1つ目の部分木をさらに[0, L), [L, R)にSplitすればよいです。区間に対する処理を行った後はMerge操作を逆に行うことで、もとの木を復元できます。
区間に対するクエリを処理するためにはノードに以下の情報を新たに保存する必要があります。
cnt
部分木のノード数です。Implicit keyを求める際に重要です。Implicit keyは上述したとおり、対応する配列の要素のインデックスのことであり、Node間の大小関係がimplicit keyの大小関係と同じであることから、あるNodeのimplicit keyは、そのNodeよりも小さいNodeの個数であり、これはcntがわかっていれば簡単に計算することができます。
すなわち、ノードVのimplicit keyはcnt(V->left)と、Vの祖先PでVがPの右の子の部分木に入っているようなVの各Pについてのcnt(P->left)の和、を合わせたものとして計算できます。
これはもっと簡単に、rootからimplicit keyを計算したい目的のノードまで順に降りていき、右に降りた場合のみcnt(V->left) + 1を累積していけば求まるので、O(logn)で計算できます。
cntの値は、それらの子がすでに更新されていると仮定すると、木DPの要領で簡単に更新できます(実際Implicit Treapではあらゆる更新は葉から根の方向に向かって行われます)。
void update_cnt(Tree t) { if (t) { t->cnt = 1 + cnt(t->l) + cnt(t->r); } }
acc
部分木のmin/sumなどを保存する変数です。この値も、それらの子がすでに更新されていると仮定すると簡単に求まります。
最小値クエリの場合のコードです(変数accの代わりにminを用いています)。
int get_min(Tree t) { return t ? t->min : INF; } void update_min(Tree t) { if (t) { t->min = min(t->value, min(get_min(t->l), get_min(t->r))); } }
lazy
部分木に対する遅延評価を保存する変数です。つまり、ノードVについて、V->value + V->lazyが、Vの実際の値となります。ここらへんは遅延Segment Treeなどと同様です。
rev
部分木の対応する区間が反転されているかのフラグです。
区間の処理の前後
Implicit Treapではあらゆる処理の前にpushdownといって、親から子方向へ遅延評価を適用していきます。
加算クエリの場合のコードです。
void pushdown(Tree t) { if (t && t->rev) { t->rev = false; swap(t->l, t->r); if (t->l) t->l->rev ^= 1; if (t->r) t->r->rev ^= 1; } if (t && t->lazy) { if (t->l) { t->l->lazy += t->lazy; t->l->min += t->lazy; } if (t->r) { t->r->lazy += t->lazy; t->r->min += t->lazy; } t->value += t->lazy; t->lazy = 0; } pushup(t); }
処理の後、pushupという関数によって親方向の情報を更新します。
void pushup(Tree t) {
update_cnt(t), update_min(t);
}
Implicit Treapの操作と実装方法
ここまでに紹介した実装がむしろ本質です。今から紹介する実装はsplit, merge, pushdown, pushupなどの組み合わせで書かれます。
split(T, key, T1, T2)
配列Tを位置keyを区切りとして2つに分割します。すなわちT1 = [0, key), T2 = [key, n)と分割されます。操作の実体はTreapのsplitと同じです。
void split(Tree t, int key, Tree& l, Tree& r) { if (!t) { l = r = nullptr; return; } pushdown(t); int implicit_key = cnt(t->l) + 1; if (key < implicit_key) { split(t->l, key, l, t->l), r = t; } else { split(t->r, key - implicit_key, t->r, r), l = t; } pushup(t); }
merge(T, T1, T2)
配列T1, T2をマージしてT2にします。Treapのmergeと同じです。
void merge(Tree& t, Tree l, Tree r) { pushdown(l); pushdown(r); if (!l || !r) { t = l ? l : r; } else if (l->priority > r->priority) { merge(l->r, l->r, r), t = l; } else { merge(r->l, l, r->l), t = r; } pushup(t); }
insert(pos, X)
配列の位置posにXを挿入します。split(T, T1, T2 pos)をして配列を[0 .. pos)と[pos, n)にsplitした後、merge(T1, T1, item)をしてXをT1にマージし、最後にmerge(T, T1, T2)で配列を1つに戻します。
void insert(Tree& t, int key, Tree item) { Tree t1, t2; split(t, key, t1, t2); merge(t1, t1, item); merge(t, t1, t2); }
erase(key)
通常のTreapのeraseと同様に実装できますが、こちらのほうがシンプルでよいと思います。要は[0, key + 1)と[ket + 1, n)に分けて、さらに前者を[0, key), [key, key + 1)に分けて、[0, key)と[key + 1, n)を最後にくっつけます。
void erase(Tree& t, int key) { Tree t1, t2, t3; split(t, key + 1, t1, t2); split(t1, key, t1, t3); merge(t, t1, t2); }
query(L, R)
配列の半開区間[L, R)についてsum/minなどを求めます。split操作で遅延評価の更新が行われるのでそのままmin/sumなどを見ればオッケーです。
最小値クエリのコードです。
int findmin(Tree t, int l, int r) { Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2, t3); int ret = t2->min; merge(t2, t2, t3); merge(t, t1, t2); return ret; }
加算クエリのコードです。
void add(Tree t, int l, int r, int x) { Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2 , t3); t2->lazy += x; t2->min += x; merge(t2, t2, t3); merge(t, t1, t2); }
reverse(L, R)
区間[L, R)を反転します。
void reverse(Tree t, int l, int r) { if (l > r) return; Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2, t3); t2->rev ^= 1; merge(t2, t2, t3); merge(t, t1, t2); }
Implicit Treapの実装
上記の機能に加え、配列のdumpやrotateを追加しています。区間クエリはmin, addに対応したものです。一般化したものは次のセクションを参考にしてください。
// for RMQ & RAQ constexpr int INF = __INT_MAX__; class ImplicitTreap { xorshift rnd; struct Node { int value, min, lazy; int priority, cnt; bool rev; Node *l, *r; Node(int value, int priority) : value(value), min(INF), lazy(0), priority(priority), cnt(1), rev(false), l(nullptr), r(nullptr) {} } *root = nullptr; using Tree = Node *; int cnt(Tree t) { return t ? t->cnt : 0; } int get_min(Tree t) { return t ? t->min : INF; } void update_cnt(Tree t) { if (t) { t->cnt = 1 + cnt(t->l) + cnt(t->r); } } void update_min(Tree t) { if (t) { t->min = min(t->value, min(get_min(t->l), get_min(t->r))); } } void pushup(Tree t) { update_cnt(t), update_min(t); } void pushdown(Tree t) { if (t && t->rev) { t->rev = false; swap(t->l, t->r); if (t->l) t->l->rev ^= 1; if (t->r) t->r->rev ^= 1; } if (t && t->lazy) { if (t->l) { t->l->lazy += t->lazy; t->l->min += t->lazy; } if (t->r) { t->r->lazy += t->lazy; t->r->min += t->lazy; } t->value += t->lazy; t->lazy = 0; } pushup(t); } void split(Tree t, int key, Tree& l, Tree& r) { if (!t) { l = r = nullptr; return; } pushdown(t); int implicit_key = cnt(t->l) + 1; if (key < implicit_key) { split(t->l, key, l, t->l), r = t; } else { split(t->r, key - implicit_key, t->r, r), l = t; } pushup(t); } void insert(Tree& t, int key, Tree item) { Tree t1, t2; split(t, key, t1, t2); merge(t1, t1, item); merge(t, t1, t2); } void merge(Tree& t, Tree l, Tree r) { pushdown(l); pushdown(r); if (!l || !r) { t = l ? l : r; } else if (l->priority > r->priority) { merge(l->r, l->r, r), t = l; } else { merge(r->l, l, r->l), t = r; } pushup(t); } void erase(Tree& t, int key) { Tree t1, t2, t3; split(t, key + 1, t1, t2); split(t1, key, t1, t3); merge(t, t1, t2); } void add(Tree t, int l, int r, int x) { Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2 , t3); t2->lazy += x; t2->min += x; merge(t2, t2, t3); merge(t, t1, t2); } int findmin(Tree t, int l, int r) { Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2, t3); int ret = t2->min; merge(t2, t2, t3); merge(t, t1, t2); return ret; } void reverse(Tree t, int l, int r) { if (l > r) return; Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2, t3); t2->rev ^= 1; merge(t2, t2, t3); merge(t, t1, t2); } // [l, r)の先頭がmになるように左シフトさせる。std::rotateと同じ仕様 void rotate(Tree t, int l, int m, int r) { reverse(t, l, r); reverse(t, l, l + r - m); reverse(t, l + r - m, r); } void dump(Tree t) { if (!t) return; pushdown(t); dump(t->l); cout << t->value << " "; dump(t->r); } public: void insert(int pos, int x) { insert(root, pos, new Node(x, rnd.random())); } void add(int l, int r, int x) { add(root, l, r, x); } int findmin(int l, int r) { return findmin(root, l, r); } void erase(int pos) { erase(root, pos); } void reverse(int l, int r) { reverse(root, l, r); } void rotate(int l, int m, int r) { rotate(root, l, m, r); } void dump() { dump(root); cout << endl; } };
モノイドを乗せる
上のコードで、minをとっている部分とlazy関連の処理をしている箇所を変更しました。ついでにインデックスアクセスとImplicit Treap上の二分探索を実装しました。
モノイドをよくわかっていないけど便利そうなのでコピペして使いたい人へ
MonoidクラスとOperatorMonoidクラスの実装は簡単だと思います。たとえばRSQ, RUQを処理したい場合は
using T = int; struct SumMonoid { static constexpr T id() { return 0; } static T op(T a, T b) { return a + b; } }; struct UpdateMonoid { static constexpr T id() { return 2e18; } static T op(T a, T b) { return b; } };
と定義すればよいです。問題はModifierですね。上の2つは使い回しが可能ですが、Modifierは区間のデータに対するクエリと区間更新に対するクエリの組み合わせによって変わります。Modifierがそもそも何を計算するかというと、accのデータが、遅延の差分(lazy)によってどう変化するかを計算します。
今の場合Modifierは以下のようになります。
struct Modifier { // lazyの結果によってaccがどう変わるか。szは部分木のサイズ static T op(T a, T b, int sz) { return b == UpdateMonoid::id() ? a : b * sz; } };
なぜなら、区間を表しているTreeがあって、その区間がbに更新されているとすると、その区間のsumは、b * 区間のサイズ(つまりTreeのサイズ)となるからです。SumMonoidやUpdateMonoidはコピペで使えますが、Modifierは毎回自分で考え直すようにしましょう。
template<class Monoid, class OperatorMonoid> class ImplicitTreap { xorshift rnd; struct Node { T value, acc, lazy; int priority, cnt; bool rev; Node *l, *r; Node(T value, int priority) : value(value), acc(Monoid::id()), lazy(OperatorMonoid::id()), priority(priority), cnt(1), rev(false), l(nullptr), r(nullptr) {} } *root = nullptr; using Tree = Node *; int cnt(Tree t) { return t ? t->cnt : 0; } T acc(Tree t) { return t ? t->acc : Monoid::id(); } void update_cnt(Tree t) { if (t) { t->cnt = 1 + cnt(t->l) + cnt(t->r); } } void update_acc(Tree t) { if (t) { t->acc = Monoid::op(acc(t->l), Monoid::op(t->value, acc(t->r))); } } void pushup(Tree t) { update_cnt(t), update_acc(t); } void pushdown(Tree t) { if (t && t->rev) { t->rev = false; swap(t->l, t->r); if (t->l) t->l->rev ^= 1; if (t->r) t->r->rev ^= 1; } if (t && t->lazy != OperatorMonoid::id()) { if (t->l) { t->l->lazy = OperatorMonoid::op(t->l->lazy, t->lazy); t->l->acc = Modifier::op(t->l->acc, t->lazy, cnt(t->l)); } if (t->r) { t->r->lazy = OperatorMonoid::op(t->r->lazy, t->lazy); t->r->acc = Modifier::op(t->r->acc, t->lazy, cnt(t->r)); } t->value = Modifier::op(t->value, t->lazy, 1); t->lazy = OperatorMonoid::id(); } pushup(t); } void split(Tree t, int key, Tree& l, Tree& r) { if (!t) { l = r = nullptr; return; } pushdown(t); int implicit_key = cnt(t->l) + 1; if (key < implicit_key) { split(t->l, key, l, t->l), r = t; } else { split(t->r, key - implicit_key, t->r, r), l = t; } pushup(t); } void insert(Tree& t, int key, Tree item) { Tree t1, t2; split(t, key, t1, t2); merge(t1, t1, item); merge(t, t1, t2); } void merge(Tree& t, Tree l, Tree r) { pushdown(l); pushdown(r); if (!l || !r) { t = l ? l : r; } else if (l->priority > r->priority) { merge(l->r, l->r, r), t = l; } else { merge(r->l, l, r->l), t = r; } pushup(t); } void erase(Tree& t, int key) { Tree t1, t2, t3; split(t, key + 1, t1, t2); split(t1, key, t1, t3); merge(t, t1, t2); } void update(Tree t, int l, int r, T x) { Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2 , t3); t2->lazy = OperatorMonoid::op(t2->lazy, x); t2->acc = Modifier::op(t2->acc, x, cnt(t2)); merge(t2, t2, t3); merge(t, t1, t2); } T query(Tree t, int l, int r) { Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2, t3); T ret = t2->acc; merge(t2, t2, t3); merge(t, t1, t2); return ret; } // [l, r)の中で左から何番目か int find(Tree t, T x, int offset, bool left = true) { if (Monoid::op(t->acc, x) == x) { return -1; } else { if (left) { if (t->l && Monoid::op(t->l->acc, x) != x) { return find(t->l, x, offset, left); } else { return (Monoid::op(t->value, x) != x) ? offset + cnt(t->l) : find(t->r, x, offset + cnt(t->l) + 1, left); } } else { if (t->r && Monoid::op(t->r->acc, x) != x) { return find(t->r, x, offset + cnt(t->l) + 1, left); } else { return (Monoid::op(t->value, x) != x) ? offset + cnt(t->l) : find(t->l, x, offset, left); } } } } void reverse(Tree t, int l, int r) { if (l > r) return; Tree t1, t2, t3; split(t, l, t1, t2); split(t2, r - l, t2, t3); t2->rev ^= 1; merge(t2, t2, t3); merge(t, t1, t2); } // [l, r)の先頭がmになるようにシフトさせる。std::rotateと同じ仕様 void rotate(Tree t, int l, int m, int r) { reverse(t, l, r); reverse(t, l, l + r - m); reverse(t, l + r - m, r); } void dump(Tree t) { if (!t) return; pushdown(t); dump(t->l); cout << t->value << " "; dump(t->r); } public: ImplicitTreap() {} ImplicitTreap(vector<T> as) { ::reverse(as.begin(), as.end()); for (T a : as) { insert(0, a); } } int size() { return cnt(root); } void insert(int pos, T x) { insert(root, pos, new Node(x, rnd.random())); } void update(int l, int r, T x) { update(root, l, r, x); } T query(int l, int r) { return query(root, l, r); } // 二分探索。[l, r)内のkでMonoid::op(tr[k], x) != xとなる最左/最右のもの。存在しない場合は-1 // たとえばMinMonoidの場合、x未満の最左/最右の要素の位置を返す int find(int l, int r, T x, bool left = true) { Tree t1, t2, t3; split(root, l, t1, t2); split(t2, r - l, t2, t3); int ret = find(t2, x, l, left); merge(t2, t2, t3); merge(root, t1, t2); return ret; } void erase(int pos) { erase(root, pos); } void reverse(int l, int r) { reverse(root, l, r); } void rotate(int l, int m, int r) { rotate(root, l, m, r); } void dump() { dump(root); cout << endl; } T operator[](int pos) { Tree t1, t2, t3; split(root, pos + 1, t1, t2); split(t1, pos, t1, t3); T ret = t3->acc; merge(t1, t1, t3); merge(root, t1, t2); return ret; } };
参考文献
Treap - Competitive Programming Algorithms
maze0123さんの実装です。 github.com
追記
update_accを演算が可換でないときにも対応できるよう修正しました。noshi91ありがとうございます。