Official

N - Product Matrix Editorial by PCTprobability


\(f(x)\) に対して \(x=1,2,4,\dots,2^M\)\(M+1\) 個の値を代入してから補間することを考えます。\(P_k = P(2^k)\) とすると、\(f(2^k) = \prod_{i=k}^{k+M-1} P_i\)\((1,1)\) 成分となります。これは、事前に \(P_{M-1},P_M\) から左右への累積積を \(\mathrm{O}(N^3M)\) かけて前計算しておくことで \(\mathrm{O}(N^3)\) で計算出来ます。

多項式補間を一般のケースにおいて \(\mathrm{O}(M \log^2 M)\) で行う方法はありますが、この問題においては TLE してしまいます。実は、評価点が等比数列である場合 \(\mathrm{O}(M \log M)\) で多項式補間を行うアルゴリズムが存在します。以降、\(v_i = f(2^i)\) と置きます。

求めたい多項式は \(f(x) = \displaystyle\sum_{i=0}^{M} v_i \prod_{j \neq i} \frac{x-2^j}{2^i-2^j}\) です。\(\displaystyle \prod_{j \neq i} (2^i - 2^j)\) は適切に式変形することで全ての \(i\) に対して \(\mathrm{O}(M)\) で列挙できます。よって、\(w_i = \displaystyle \frac{v_i}{\prod_{j \neq i} (2^i - 2^j)}\) と定義することによって問題を \(\displaystyle \sum_{i=0}^{M} w_i \prod_{j \neq i} (x - 2^j) = \left(\prod_{i=0}^{M} (2^i - x)\right)\left(\sum_{i=0}^{M} \frac{w_i}{2^i - x}\right)\) を求める問題に帰着されます。

前者は \(\displaystyle \prod_{i=0}^{k} (2^i - x)\) から \(\displaystyle \prod_{i=0}^{2k+1} (2^i - x)\)\(\mathrm{O}(k \log k)\) で求められるため \(\mathrm{O}(M \log M)\) で求められます。(Cauchy Binomial Theorem を使い \(\mathrm{O}(M)\) で求めることも出来ます。)

後者は \(\displaystyle \sum_{i=0}^{M} \frac{w_i}{2^i - x} = \sum_{i=0}^{M} \frac{\frac{w_i}{2^i}}{1 - \frac{x}{2^i}} = \sum_{i=0}^{M} \frac{w_i}{2^i}\sum_{j=0}^{\infty} \frac{x^j}{2^{ij}} = \sum_{j=0}^{\infty} \sum_{i=0}^{M} \frac{w_i}{2^{(i+1)j}} x^j\) と式変形を行うことで、\(g(x) = \displaystyle \sum_{i=0}^{M} w_i x^{i+1}\) として \(g(2^{-k})\)\(k=0,1,\dots,M\) で求める問題に帰着されます。

\(t_i = \frac{i(i-1)}{2}\) と置くと、\(g(2^{-k}) = \displaystyle \sum_{i=0}^{M+1} g_i 2^{-ik} = \sum_{i=0}^{M+1} g_i 2^{-t_{i+k}}2^{t_i}2^{t_k}\) と式変形をすることが出来ます。すると、\(i,i+k,k\) だけの項に分離することが出来るため畳み込みを用いて \(\mathrm{O}(M \log M)\)\(g(2^{-k})\) を列挙することが出来ます。(このパートのことは chirp z-transform と呼ばれています。)

途中で \(\bmod\ 10^9+7\) で畳み込みを行う必要がありますが中国剰余定理等を用いることで \(\bmod\ 10^9+7\) でも長さ \(k\) の畳み込みを \(\mathrm{O}(k \log k)\) で行えます。

よって、この問題を \(\mathrm{O}(N^3M + M \log M)\) で解くことが出来ます。

参考:https://noshi91.github.io/algorithm-encyclopedia/polynomial-interpolation-geometric

余談

この問題でかなり厳しい TL になってしまったことや \(\bmod\ 10^9+7\) を要求した事情を説明します。この問題はダブリングを用いることで \(M=k\) のケースから \(M=2k\) のケースの答えを \(\mathrm{O}(N^3 k \log k)\) で得ることが出来るため、全体で \(\mathrm{O}(N^3 M \log M)\) でこの問題を解くことが出来ます。

更に、一般的に \(N \times N\) で各要素が \(M\) 次の多項式であるような行列の積は素直にやると \(\mathrm{O}(N^3 M \log M)\) かかりますが、FFT の結果を使いまわし IFFT をまとめて行うことで \(\mathrm{O}(N^2 M \log M + N^3 M)\) で行うことが出来ます。よって、この高速化を上記と組み合わせるとこの問題を \(\mathrm{O}(N^2 M \log M + N^3 M)\) で解くことが出来ます。この解法と想定解を区別するために非常に厳しい制約と TL 設定に加え \(\bmod\ 10^9+7\) を要求することになってしまいました。この調整によって定数倍等で落ちてしまった人や \(\mathrm{O}(N^2 M \log M + N^3 M)\) が想定解だと思い通らなかった人等には申し訳なく思っております。有志コンで上記のような面白いアルゴリズムを布教することが目的ということで、どうかご勘弁ください。

posted:
last update: