E - Least Elements Editorial by rsk0315

セグ木を用いた実装

\(j\in\{i, i+1, \dots, i+M-1\}\) を、\(A_j\) の値によって小さい方から \(K\) 個選んだものを \(L_i\)、残りを \(R_i = \{i, i+1, \dots, i+M-1\}\setminus L_i\) とします。 すなわち、

  • \(L_i \cup R_i = \{i, i+1, \dots, i+M-1\}\)
  • \(|L_i| = K\)
  • \(\max\, \{A_j \mid j\in L_i\} \le \min\, \{A_j \mid j\in R_i\}\)

とします。\(\max\, \{A_j \mid j \in L_i\} = \min\, \{A_j\mid j\in R_i\}\) となるときの分割の仕方は任意ですが、実装上は \(j\) の値が小さい方を \(L_i\) に入れるのが楽でしょう。

さて、4 本のセグ木を用います。

  • sum0: \(j\in L_i\) のとき \(\text{\texttt{sum0}}[j] = 1\)、そうでないとき \(\text{\texttt{sum0}}[j] = 0\)。区間和を管理する。
  • sum1: \(j\in L_i\) のとき \(\text{\texttt{sum1}}[j] = A_i\)、そうでないとき \(\text{\texttt{sum1}}[j] = 0\)。区間和を管理する。
  • used: \(j\in L_i\) のとき \(\text{\texttt{used}}[j] = (A_i, i)\)、そうでないとき \(\text{\texttt{used}}[j] = (-\infty, i)\)。区間最大値を管理する。
  • unused: \(j\in R_i\) のとき \(\text{\texttt{unused}}[j] = (A_i, i)\)、そうでないとき \(\text{\texttt{unused}}[j] = (\infty, i)\)。区間最小値を管理する。

\((L_i, R_i)\) から \((L_{i+1}, R_{i+1})\) を求める際には、\(i\notin L_{i+1}\) なので、\(\text{\texttt{sum0}}[i] \gets 0\) などとします。\(i+M \in R_{i+1}\) なので、\(\text{\texttt{sum1}}[i+M] \gets 1\) などとします。

また、 used における最大値を求めることで、どれを \(R_{i+1}\) に移すべきか(あるいは unused の最小値によって、どれを \(L_{i+1}\) に移すべきか)がわかります。

\(M = K\) のケースは別途処理すると楽になるかもしれません。

提出 #37192377


上記のようにセグ木を用いて優先度つきキューを実装すると、添字などに関連するクエリも処理できるようになります。覚えておくと「普通のヒープだと○○のせいでうまくいかないな」となったときに対処しやすくなるかもしれません。

関連:https://rsk0315.hatenablog.com/entry/2022/01/13/233015

posted:
last update: