ABC 440 に参加しました。
D問題Forbidden List 2について、自分の解法が公式解説とは違う方針だったので、記事にします。
私も最初は公式解説と同じように二分探索で頑張ろうとしましたのですが、端数合わせに苦労したので、方針変更しました。
問題概要
長さ$L$の整数列に対し、$X$以上の整数のうち、小さいほうから$Y$番目のものを求めよ。
ただし$X,Y$のクエリはたくさん来るので、一回のクエリあたり、$O(log N)$で解いてね。
方針
使用可能な数だけ管理しよう。
もう少し詳しい考察
使用可能な数のうち、数が連続している部分ごとにグループ化し、グループの最小値について、何番目かを覚えておけば良さそう。(ここで、数が連続しているというのは、差分が$1$という意味です)
例えば$A_i=[1,2,5,8,10]$の時を考えます。
使用できる数は$3,4,6,7,9,11,12,...$となりますので、連続部分ごとにグループ化すると、$[3,4],[6,7],[9],[11,12,...]$となります。
ですので、下記のような状態になります。
| index | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|---|---|---|---|---|---|---|---|
| 値 | 3 | 4 | 6 | 7 | 9 | 11 | 12 |
| グループの最小値？ | YES | NO | YES | NO | YES | YES | NO |
というわけで、$(1,3),(3,6),(5,9),(6,11)$という値をあらかじめ求めておくことで、「小さい方から$Y$番目の値は？」というクエリに対し、$O(\log_{2} N)$で回答できます。
(グループの数が$O(N)$になるため、どのグループに含まれるかの探索に$O(\log_{2} N)$、そこからの計算に$O(1)$。)
なお、この値はBTreeMapで管理するのが楽だと思います。
なお、$X$未満の数のうち使用可能な数の個数を$Z$とすると、使用可能な数のうち$Y+Z$番目に小さい数を求めれば良いことになります。
また$Z$はbinary_search等により$O(\log N)$で求められますので、各クエリを$O(\log N)$で求められます。
提出コード
余談
二分探索した後の端数処理が苦手です。
Fも公式解説とは違い、セグ木なしで解いたので記事にする予定。