K - Kth Sum Editorial
by
SSRS
以下、\(A_i, B_j, C_k\) の最大値を \(M\) とします。また、\(A,B,C\) は昇順にソートされているとします。
答えについて二分探索することを考えると、以下の問題を \(\mathrm{O}(\log M)\) 回解くことに帰着されます。
\(A_i + B_j + C_k \leq x\) なる \((i,j,k)\) の個数は \(K\) 個以下であるか?
\(i\) を固定したとき \(A_i + B_j + C_k \leq x\) となる \((j,k)\) の個数を \(f(i)\) とおくと、それぞれの \(f(i)\) は \(B,C\) の単調性を用いて \(\mathrm{O}(N)\) 時間で求められます。このとき、\(\displaystyle\sum_{i=1}^N f(i) \leq K\) かどうかを判定すればよいです。
\(A\) は広義単調増加なので、\(f(i)\) は \(i\) について広義単調減少となります。
最初に、\(f\left(\sqrt{\frac KN}\right)\) を求めます。(\(\sqrt{K/N}\) が整数でない場合は、整数になるように適切に丸めるものとします。以下、整数でない場合の議論は省略します。) この値が \(\sqrt{NK}\) より大きい場合は、\(f(1)+f(2)+\dots+f\left(\sqrt{\frac KN}\right) > \sqrt{\frac KN} \cdot \sqrt{NK} = K\) なので、全体の和も \(K\) より大きいことがわかります。
\(f\left(\sqrt{\frac KN}\right) < \sqrt{NK}\) である場合は、以下のようにして \(\displaystyle\sum_{i=1}^N f(i)\) を求められます。
- \(f(1), f(2), \dots, f\left(\sqrt{\frac KN}-1\right)\) はそれぞれ \(\mathrm{O}(N)\) 時間で求められる。
- 最初に、\(B_j + C_k\) の下位 \(\sqrt{NK}\) 個の値を前計算しておく。前計算は優先度付きキューを用いると \(\mathrm{O}(\sqrt{NK} \log N)\) 時間でできる。\(f\left(\sqrt{\frac KN}+1\right), \dots, f(N)\) はすべて \(\sqrt{NK}\) 未満であるため、二分探索でそれぞれ \(\mathrm{O}(\log \sqrt{NK})\) 時間で求められる。
全体の時間計算量は \(\mathrm{O}(\sqrt{NK}(\log N + \log M) + N \log N \log M)\) となります。
posted:
last update: