この記事は レジリエントでウェルビーイングなぴょこりんクラスタ Advent Calendar 2024 の為に書いたものです。
はじめに
この記事は前回の続きです。
続きものとなりますので、前回の内容はこちらから参照してください。
Treap と Implicit Treap というものがあり、今回使うのは後者。
Implicit Treap は以下の操作を で行えるデータ構造、という話をしました。。
- ある位置 p に要素 x を挿入する
- ある位置 p にある要素を削除する
- ある範囲の要素の最大値を計算する
- ある範囲の要素の最小値を計算する
- ある範囲の要素すべての総和を計算する
- ある範囲の要素すべてに値 a を加算する
- ある範囲の要素を左右反転させる
今回は、このデータ構造を使って改めて MetaHackerCup 2024 PracticeRound D2 問題を解いていきましょう。
Treap をつかって問題を解いてみる
期間が空いて忘れてしまったかもしれませんが、
今回解きたかった問題の概要は1日目の記事を参照してください。
Treap を使って解けるよ、と言っていた公式 (原文) が言っていることを改めておさらいすると、
力 でストーンを投げた場合、
ストーンが止まるまでには他のストーンに衝突が起こったりもするけれど、
最終的な石が止まる位置というのは、
ストーンがない座標のみを左から数えて 番目の位置である。
また、それ以前に存在していたストーンは、すべて座標が -1 される。
つまりは、特定の P に値を入れることと、P以前の値を1引くという操作を高速に（データ量に触れていませんが、この処理を で）計算できれば良い。
こういうのに便利なデータ構造として Treap ってのがあるので、それをうまく使えば計算できるよ。
というものでした。
解説をもとに Treap の解法を書いてみようとする、が…
P 以前の値を１引くことも P に値を入れることも Treap でできるから楽勝…って思うかもしれませんが、実はここに一山あります。
特定のPに値を入れることと、P以前の値から1を引くという操作を Treap で行うことを踏まえてコードを書いてみると、以下のようになるかと思います。
int Case;
cin >> Case;
int ans = 0;
for (int ci = 0; ci < Case; ci++) {
int N, G;
cin >> N >> G;
ImplicitTreap t;
rep(i, N) {
int E;
cin >> E;
// determine E-th empty space as p, and index of p in ImplicitTreap
//
// ???
//
t.insert(p_index, p);
t.add(0, p_index, -1);
}
// find the position nearest to G as Q, and its index
//
// ???
//
cout << "Case #" << ci + 1 << ": " << N - q_index << " " << abs(Q - G)
<< endl;
}
Implicit Treap に値を入れたり範囲加算・減算をしたりするときにはその index が必要だし、
そもそも E 番目のストーンのない場所である p を計算する方法はいまいちピンときませんね。
また、その山場を突破したとしても、Gに最も近いストーンの位置を求めることについても、
Treap が提供する操作では直に実現できなさそうです。
僕はこれが分からずまたしばらく思い悩んだのですが、
しばらく考えた結果（これが最適化はわかりませんが）一応の答えにはたどり着くことができました。
結論としては、どちらも二分探索をつかいました。
p（ストーンがない E 番目の座標）や p_index (Implicit Treap 上の p の位置）の求め方
ある力Eで発射したストーンが最終的に到達する位置 p は、
ストーンがない E 番目の座標として計算できるわけですが、
これはすなわち、p_index が計算できるのであれば、 p = E+p_index と計算することができるため、
p_index を計算することができれば十分です。
Implicit Treap t にはストーンの現在の座標が昇順に入っているわけなので、
ストーン t[i] の位置エネルギーを Energy(t[i]) と表記すると、
Energy(t[left]) <= E < Energy(t[right]) となるように left, right を二分探索で更新すれば、
最終的にPが挿入されるべき個所（p_index) を決めることができます。
// determine E-th empty space as p, and index of p in ImplicitTreap int left = -1; int right = i; while(left+1<right){ int mid = (left+right)/2; // get t[mid] int mid_elem = t.getMax(0,mid+1); // calculate energy of t[mid] stone int mid_energy = mid_elem - mid; if( mid_energy <= E) { left = mid; } else{ right = mid; } } int p_index = right; int p = E+p_index;
ちなみに上記のコードでも //get t[mid] 付近で使っていますが、
t[i] の座標は t.getMax(0, i+1) で計算することが可能です。
一般的な Implicit Treap では要素が昇順に並んでいるとは限らないのですが、
今回は要素を昇順に並べているため、特別な関数を増やす必要もないんですね。
もっとも、 i 番目の要素参照なんて昇順でなくても使いたいものなので、 Treap の実装に追加していてもよい気はしますけど。
Q（Gに最も近いストーンの座標）の求め方
こちらについてはもっと単純で、 t[left] <= G < t[right] が満たされるように left, right を二分探索で更新すれば、 Gの左右にあるストーンが求まるので、それぞれ検証してみて値の近い方を選択すればOKとなります。
// find the position nearest to G as Q, and its index int left = -1; int right = N; while(left+1 < right){ int mid = (left+right)/2; // get t[mid] int mid_elem = t.getMax(0,mid+1); if(mid_elem <= G){ left = mid; } else{ right = mid; } } int cand1 = left; int stone1 = t.getMax(0,cand1+1); // t[cand1]; int cand2 = right; int stone2 = t.getMax(0,cand2+1); // t[cand2]; int q_index = cand1; int Q = stone1; if(abs(stone1-G) >= abs(stone2-G) && cand2 != N){ // t[cand2] = t.getMax(...) を使っていると cand2 =N のケースがコーナーケースになるため、最後の条件ではじく q_index = cand2; Q = stone2; }
最終的なコードと動作確認
ここまでの話をまとめたコード
typedef struct Node { int val; int priority; int cnt; int min; int max; int acc; ll sum; Node* left; Node* right; Node(int val, int priority) : val(val), priority(priority), cnt(1), sum(val), min(val), max(val), acc(0), left(nullptr), right(nullptr) {} } Node; class ImplicitTreap { private: Node* root; random_device seed_gen; mt19937 engine{seed_gen()}; uniform_int_distribution<int> distribution{INT_MIN, INT_MAX}; int count(Node* t) { return !t ? 0 : t->cnt; } ll sum(Node* t) { return !t ? 0 : t->sum; } int getMax(Node* t) { return !t ? INT_MIN : t->max; } int getMin(Node* t) { return !t ? INT_MAX : t->min; } // update information of t using information of its children // this function assumes all the information of the child are updated Node* update(Node* t) { if (t) { t->cnt = count(t->left) + count(t->right) + 1; t->sum = sum(t->left) + sum(t->right) + t->val; t->min = min(t->val, min(getMin(t->left), getMin(t->right))); t->max = max(t->val, max(getMax(t->left), getMax(t->right))); } return t; } // push acc, lazy info of node t into its child void push(Node* t) { if (t) { if (t->left) { t->left->acc += t->acc; t->left->max = t->left->max + t->acc; t->left->min = t->left->min + t->acc; t->left->sum = t->left->sum + (t->left->cnt) * t->acc; } if (t->right) { t->right->acc += t->acc; t->right->max = t->right->max + t->acc; t->right->min = t->right->min + t->acc; t->right->sum = t->right->sum + (t->right->cnt) * t->acc; } t->val = t->val + t->acc; t->acc = 0; } } // merge Node* l and Node *r into unique tree Node* merge(Node* l, Node* r) { if (l) push(l); if (r) push(r); if (!l || !r) return l ? l : r; if (l->priority > r->priority) { l->right = merge(l->right, r); return update(l); } else { r->left = merge(l, r->left); return update(r); } } // split Node* into l=[0,k), r=[k,n) pair<Node*, Node*> split(Node* t, int k) { if (!t) return make_pair(nullptr, nullptr); push(t); if (count(t->left) >= k) { pair<Node*, Node*> s = split(t->left, k); t->left = s.second; return make_pair(s.first, update(t)); } else { pair<Node*, Node*> s = split(t->right, k - count(t->left) - 1); t->right = s.first; return make_pair(update(t), s.second); } } // insert item into pos of t void insert(Node*& t, int pos, Node* item) { if (!t) { t = item; return; } pair<Node*, Node*> s = split(t, pos); t = merge(merge(s.first, item), s.second); } // erase t[pos] void erase(Node*& t, int pos) { if (!t) return; pair<Node*, Node*> p = split(t, pos + 1); pair<Node*, Node*> p2 = split(p.first, pos); if (p2.second) delete p2.second; t = merge(p2.first, p.second); } void print(Node* t, int depth) { if (!t) { for (int i = 0; i < depth; i++) cout << " "; cout << "null" << endl; return; } else { print(t->right, depth + 1); for (int i = 0; i < depth; i++) cout << " "; cout << "val: " << t->val << ", " << "priority: " << t->priority << ", " << "count: " << t->cnt << ", " << "sum: " << t->sum << ", " << "max: " << t->max << ", " << "min: " << t->min << ", " << "acc: " << t->acc << endl; print(t->left, depth + 1); } } // return the sum of [left,right) ll sum(Node*& t, int left, int right) { pair<Node*, Node*> p = split(t, right); pair<Node*, Node*> p2 = split(p.first, left); ll ans = sum(p2.second); t = merge(merge(p2.first, p2.second), p.second); return ans; } // return the min of [left,right) ll getMin(Node*& t, int left, int right) { pair<Node*, Node*> p = split(t, right); pair<Node*, Node*> p2 = split(p.first, left); ll ans = getMin(p2.second); t = merge(merge(p2.first, p2.second), p.second); return ans; } // return the max of [left,right) ll getMax(Node*& t, int left, int right) { pair<Node*, Node*> p = split(t, right); pair<Node*, Node*> p2 = split(p.first, left); ll ans = getMax(p2.second); t = merge(merge(p2.first, p2.second), p.second); return ans; } // add value to the all elements in [left,right) void add(Node*& t, int left, int right, int value) { pair<Node*, Node*> p = split(t, right); pair<Node*, Node*> p2 = split(p.first, left); if(p2.second) p2.second->acc += value; t = merge(merge(p2.first, p2.second), p.second); } public: ImplicitTreap() { root = nullptr; } // insert val into ImplicitTreap[pos] void insert(int pos, int val) { insert(root, pos, new Node(val, distribution(engine))); } // erase ImplicitTreap[pos] void erase(int pos) { erase(root, pos); }; // sum ImplicitTreap[left,right) ll sum(int left, int right) { return sum(root, left, right); } // max ImplicitTreap[left,right) int getMax(int left, int right) { return getMax(root, left, right); } // min ImplicitTreap[left,right) int getMin(int left, int right) { return getMin(root, left, right); } // add value to all of ImplicitTreap[left,right) void add(int left, int right, int value) { add(root, left, right, value); } // print void print() { print(root, 0); } }; int main() { int Case; cin >> Case; int ans = 0; for (int ci = 0; ci < Case; ci++) { int N, G; cin >> N >> G; ImplicitTreap t; rep(i, N) { int E; cin >> E; // determine E-th empty space as p, and index of p in ImplicitTreap int left = -1; int right = i; while(left+1<right){ int mid = (left+right)/2; // get t[mid] int mid_elem = t.getMax(0,mid+1); // calculate energy of t[mid] stone int mid_energy = mid_elem - mid; if( mid_energy <= E) { left = mid; } else{ right = mid; } } int p_index = right; int p = E+p_index; t.insert(p_index, p); t.add(0, p_index, -1); } // find the position nearest to G as Q, and its index int left = -1; int right = N; while(left+1 < right){ int mid = (left+right)/2; // get t[mid] int mid_elem = t.getMax(0,mid+1); if(mid_elem <= G){ left = mid; } else{ right = mid; } } int cand1 = left; int stone1 = t.getMax(0,cand1+1); // t[cand1]; int cand2 = right; int stone2 = t.getMax(0,cand2+1); // t[cand2]; int q_index = cand1; int Q = stone1; if(abs(stone1-G) >= abs(stone2-G) && cand2 != N){ q_index = cand2; Q = stone2; } cout << "Case #" << ci + 1 << ": " << N - q_index << " " << abs(Q - G) << endl; } return 0; }
使って公式のテスト入力を実行し、出力が以前に解いた解法の出力と一致することが確認できました。
ちなみに計算量は Treap 操作と二分探索を要素数だけ繰り返すので とかだと思います。
なんとなく実行時間は Implicit Treap の方が遅かった気がする。
まぁでもコンテストの実行制限には（たぶん）引っかからない速度感だとおもいます。
おわりに
というわけで、今回は四回にわたって、Treap がつかえるといわれる問題の Treap 不要解と、 Treap 利用解を説明してきました。
Treap 解の実装までひたすらかかった挙句に不要解より遅い気がするし実装も複雑だし大変じゃね？
... なんて思った人も多いと思うのですが、
実のところ、個人の体感として Treap 関係の問題はこういった現象が起きてることがほとんどなんですよね。
だからこそ、１０年以上競技プログラミングをやっている自分も Treap の実装を持ってなかったし、 Implicit Treap と Treap の違いを把握もしていませんでした。
（というか、調べても良くわからないで挫折…を繰り返していた）
そんな背景があったので、かなり読者置いてけぼりの記事だった自覚はあれど、
アドベントカレンダーという体で重い腰を上げてちゃんと向き合うことができて個人的にはよかったです。満足。
おまけ
Q ところで Meta Hacker Cup の Practice Round についてばかり話してたけど、肝心の本戦はどうだったの？
A 寝落ちにより参加できませんでした。
