ABC170 E - Smart Infantsのmultisetを利用しない解き方

AtCoder Beginner ContestのE問題でmultisetを用いた解説が書いてありましたがpythonなどmutisetが用意されていない言語などもあると思うのでPriority Queue(heapq)を用いた代替アルゴリズムを紹介したいと思います。
問題自体の解説はしないのでそこは公式の解説を読んでください。

問題

ABC170 E - Smart Infants
https://atcoder.jp/contests/abc170/tasks/abc170_e

前提知識

multiset

内部的には二分探索木が使われています。これによりデータの挿入で毎回O(logN)かかる代わりに常にソートされている状態が維持されています。そのため、削除や検索もO(logN)で行うことが可能です。
自分で二分探索木を実装しておくことで標準で提供されていない言語でも使用することはできますが平衡を保つのには工夫が必要です。(工夫例:【木マスター養成講座】7-1. Splay木ってなに〜？説明編【競プロかつっぱ】)

Priority Queue

GCCでは内部的に二分ヒープが使用されています。
pythonのheapqでも内部的には同じデータ構造が用いられています。
データの挿入は二分探索木と同じく、毎回O(logN)かかります。ヒープ木なので内部はソートされておらず、検索や最大(最小)値以外の削除は高速にはできませんが、常に最大(最小)値を計算量O(1)で取得することはできます。二分探索木より勝っている点は、特に工夫をしなくても常に平衡が保たれているため計算量が安定するという点です。

解法

要点

今回の問題では

1. 各幼稚園においてレートの最大値を取ってくること
2. レートの最大値の集合の最小値を取ってくること
3. 入園するときに値をそれぞれの集合へ挿入すること
4. 転園の際に値をそれぞれの集合から削除すること

が常に高速でできれば良いことになります。

公式の解説でも述べられているようにmultisetを用いることでそれぞれの処理は高速に行うことができますが、Priority Queueでは4番目の処理を高速で行うことができないため発想を変える必要があります。

今回の問題を解くに当たって、4番目の条件は、1,2番目の条件の邪魔にならなければ必ずしも行わなければならない処理ではないことに着目しましょう。
つまり、実際にデータの中から削除されていなくても各幼稚園におけるレートの集合の最大値やレートの最大値の集合の最小値に影響が無い場合は構わないわけです。

具体的なアルゴリズム

1,2番目の条件で実装は変わらない(最大最小がひっくり返るぐらいな)ので1番目の条件(集合の最大値の取得)を例に説明します。

ここで、Priority Queueを2本持つことを考えます。

1. 集合全体を管理するもの
2. 1のうち削除されたはずの要素を保持するもの

そして入園した時には1に、転園したときは2に要素を挿入します。1と2の最大が同じ値の場合、両方から削除します。
これにより、この幼稚園の最大値は1の最大値と常に同じになります。

なんとなくのコードを載せておきますがstdやincludeなど端折っているので環境によってはコピペしても動かないです。

コード.cpp

priority_queue<int> in;
priority_queue<int> out;

// 入園するときの処理
void infant_in(int i){
    in.push(i);
}

// 転園するときの処理
void infant_out(int i){
    out.push(i);
    while(!out.empty() && in.top() == out.top()){
        in.pop();
        out.pop();
    }
}

// 最大値の取得
int get_max(){
    return in.top();
}

具体例を見ましょう。幼稚園にレートが3,5,8の園児が入園してきてその後レートが5,8の園児がこの順に転園していくことを考えます。

> 3人の園児たちが入園
実際の幼稚園 = {3,5,8}
in = [3,5,8] out = []   //この時点での最大値はもちろん8

> レートが5の園児が転園
実際の幼稚園 = {3,8}
in = [3,5,8] out = [5]  //最大値は8で合っている

> レートが8の園児が転園
実際の幼稚園 = {3}
in = [3,5,8] out = [5,8]
in = [3,5]   out = [5]
in = [3]     out = []   //最大値は3で合っている

こんな感じで、最大値を取得する上では問題ないことがわかります。

計算量

気になる計算量なのですが、今回の問題においてネックになりそうなinとoutから値を削除する回数は最大Q回なので単独でボトルネックになることはなく、計算量は値の挿入時にかかるO((N+Q)logN)になります。
これは十分に高速と言えます。
(6/17 計算量が正しくなかったので修正しました)