Official

L - K番目の絶対値/K-th Abs Editorial by Nyaan


まず、ナイーブなアルゴリズムとして全ての部分列に対してスコアを計算して \(K\) 番目を求める方法が考えられますが、部分列は \(\mathrm{O}(N^2)\) 個あるので TLE してしまいます。

この問題のように \(K\) 番目の要素を求める問題では、二分探索 を利用したアルゴリズムが有効になります。
答えを \(a\) とおくと、スコアが \(K\) 以上の部分列の個数と \(a\) の間に以下の関係が成り立つことがわかります。

  • \(x \lt a\) のとき、スコアが \(x\) 以下の部分列は \(K\) 個未満
  • \(x \geq a\) のとき、スコアが \(x\) 以下の部分列は \(K\) 個以上

よって、各 \(x\) に対して 「スコア \(x\) 以下の部分列は \(K\) 個以上あるか?」という判定問題を十分高速に解くことができれば、二分探索を利用して \(\mathrm{O}(\log M)\) 回判定問題を解くことでこの問題を解くことができます。 ( \(M\) はスコアの最大値と最小値の差)
スコアの最大値は \(N \max(A) = 3 \times 10^{14}\) 、最小値は \(0\) なので、二分探索は高々 \(\lceil \log_2 (3 \times 10^{14} ) \rceil = 49\) 回で済みます。

「スコア \(x\) 以下の部分列は \(K\) 個以上あるか?」という判定問題を高速に解く方法を考えます。
数列 \(A_i\) \((1 \leq i \leq N)\) に対して、数列 \(R_i\) \((0 \leq i \leq N)\)

\[R_0 = 0, R_i = R_{i-1} + A_i \ (1 \leq i \leq N)\]

を満たすように定めます。このとき \(0 \leq i \lt j \leq N\) を満たす整数の組 \((i, j)\) に対して \(R_j - R_i\) の値が部分列 \((i, j \rbrack\) の要素の総和と一致します。よって判定問題は \(R\) を使うと次のように言い換えられます。

  • \(0 \leq i \lt j \leq N\) を満たす整数の組 \((i, j)\) であって、\(| R_j - R_i | \leq x \) を満たすものは \(K\) 個以上存在するか?

ここで、 \( | R_j - R_i | \leq x\) ならば \(|R_i - R_j| \leq x\) である点に着目すると、 \(i \lt j\) という順序関係を取っ払って \(i \neq j\) に変えた場合は条件を満たすペアの個数は 2 倍になります。すなわち、判定問題は次の命題に言い換えられます。

  • \(0 \leq i, j \leq N , i \neq j\) を満たす整数の組 \((i, j)\) であって、\(| R_j - R_i | \leq x \) を満たすものは \(2 \times K\) 個以上存在するか?

この問題は \(R\) をあらかじめソートしたあとに尺取り法と呼ばれるアルゴリズムで \(\mathrm{O}(N)\) で計算できます。よって判定問題を \(\mathrm{O}(N)\) で解くことができました。

以上よりこの問題を \(\mathrm{O}(N \log M)\) ( ただし \(M = N \max(A)\) )でこの問題を解くことができました。

posted:
last update: