集合 \({1,2,...,N}\) (\(N\) は巨大) から数を消していきながら、現在 \(k\) 番目の数を求めるテク

突然ですが、次の問題を解きたいです。(これは Move the Coins - Creating Tests の部分問題です。全部やるとむつごいので次の問題だけ取り出します。)

\(1,2,...,N\) をそれぞれ \(1\) 個ずつ含む集合があります。以下の \(Q\) 個のクエリを処理してください。

1 x : 集合から \(x\) を削除する。
2 k : 集合にある数の中で、 \(k\) 番目に小さいものを求める。

\(N \le 10^{18} , Q \le 10^5\)

 もちろん、集合が巨大すぎるのでBITなりsetなりで管理するのは困難です。必要な分だけ作るセグ木なら出来そうだけど非常に面倒…という感じでしょうか。

そこで活躍するのが treeくん。 詳しくは Policy Based Data Structures - 競プロ練習記録 を参照してください。

treeは、setの機能に加えて「 \(k\) 未満の要素がいくつあるか?」 というクエリにも対応可能です。

すると、これを利用して、(\(x\)以下の今残っている要素の数) = \(x\) - ( \(x+1\) 未満の既に消された要素の数) という式を実装できます。ここまで来れば、あとは二分探索により \(k\) 番目の数が何か求められます。

計算量は、クエリあたり \(O( \log^2 Q ) \) 程度と、√より少し軽いか同じくらいでしょうか。額面では \(Q \le 2 \times 10^5\) あたりまでは安全圏そうです。 しかし、この解法の強みはシンプルな実装(実装例を見て頂くと分かります)に加え、何より \(N\) に依存しないことでしょう。空間計算量もtreeの \(O(Q)\) で済みます。

C++での実装例は こちら (この例では、元の集合のサイズが無限に大きくても、オーバーフローしない限り対応可能です)

ありそうで意外と見かけない気がしたので、記事化しました〜

ここから追記

平衡二分探索木方向でうまいことやるなりすれば \(O(\log Q)\) まで落ちます(他にも落とせる方法はあるっぽいです) とっさに書くなら間違えにくい(と個人的には思う)上の実装でも充分だと思いますが、計算量で困った場合には改良しましょう

よくよく考えればC++にmapとかいうのもありましたね、あれでBIT(Fenwick Tree)やっても楽にできそう?