C - Sum of Three Integers Editorial
by
Nyaan
問題を解くために補助的な配列 \(f(n), g(n)\) を次のように定義して計算します。
\(f(n)\) を \(A_i = n\) である \(i\) の個数として定義します。(すなわち \(f(n)\) は頻度列です) \(f(1), f(2), \dots, f(X)\) は \(\mathrm{O}(N + X)\) で計算できます。
\(g(n)\) を \(A_i + A_j = n, i \lt j\) を満たす整数の組 \((i, j)\) の個数として定義します。\(f(n)\) と \(g(n)\) の間には
\[g(n) = \frac{\left(\sum_{i=1}^{n-1} f(i) f(n-i)\right) - \left(f\left( \frac{n}{2}\right) \text{ if } n \bmod 2 = 0 \text{ else } 0\right)}{2}\]
という関係式が成り立ちます。よって \(g(1), g(2), \dots, g(X)\) は畳み込みで \(\mathrm{O}(X \log X)\) で計算できます。
- 畳み込み後の要素の値は最大で \(10^{12}\) になるため通常の \(\text{mod }998244353\) 畳み込みだけでは上手くいかない点に注意が要ります。C++ では ACL の
convolution_ll
関数を用いるのが簡明です。convolution_ll
が存在しない環境でも、十分長い FFT が可能な \(10^8\) 以上の素数 mod を 2 種類選んできてそれぞれ畳み込んだ後に Garner’s algorithm(中国剰余定理, CRT) で復元すれば求めたいものを計算できます。また、複素数 FFT を用いた畳み込みでも与えられたテストケースに対して正確に計算できることを確認していますが、誤差に関する適切な評価はできていません。
さて、次の条件を満たす整数 \(s\) \((1 \leq s \leq N)\) を探す問題を考えます。
- 整数の組 \((t, u)\) \((1 \leq t \lt u \leq N)\) であって次の条件を全て満たすものが存在する。
- \(s \neq t, s \neq u\)
- \(A_s + A_t + A_u = X\)
条件を満たす \(s\) を発見できれば、問題の答えとなる \((i, j, k)\) を発見することは \(f(n)\) を用いて \(\mathrm{O}(N + X)\) で行うことが出来ます(詳細は略します) 。また、条件を満たす \(s\) が存在しなければ答えは解なしです。よって、条件を満たす \(s\) を発見する問題が解ければよいです。
\(s\) を \(1 \leq s \leq N\) の範囲で全探索します。\(s\) を固定した時に、\(1 \leq t \lt u \leq N, s \neq t, s \neq u, A_s + A_t + A_u = X\) を満たす整数の組 \((t, u)\) の個数は \(f(n), g(n)\) を用いて
\[g(X - A_s) - \left( f(X - 2 A_s) - (1 \text{ if } X=3 A_s \text{ else } 0 )\right)\]
と表すことができます。(詳細は略します。またここで \(n \leq 0\) に対して \(f(n) = 0\) とします。) よって上式を計算すれば \(s\) が条件を満たすかどうかは \(\mathrm{O}(1)\) で判定できるため、条件を満たす \(s\) の発見は \(\mathrm{O}(N)\) で行うことが出来ます。
よってこの問題を全体で \(\mathrm{O}(N + X \log X)\) で解けて、これは十分高速です。
posted:
last update: