中国のブログ[1]に面白いアルゴリズムがあったので、紹介します。
問題
次のような多変数畳み込み（切り捨て）を計算するアルゴリズムを構築します：
$$f(x_1,x_2,\ldots,x_K)=\sum a_{i_1,i_2,\ldots,i_K}x_1^{i_1}x_2^{i_2}\cdots x_K^{i_K},$$$$g(x_1,x_2,\ldots,x_K)=\sum b_{i_1,i_2,\ldots,i_K}x_1^{i_1}x_2^{i_2}\cdots x_K^{i_K}$$ が与えられます。\(f,g\) の積を途中で切り捨てた$$f(x_1,x_2,\ldots,x_K)g(x_1,x_2,\ldots,x_K) \bmod (x_1^{n_1},x_2^{n_2},\ldots,x_K^{n_K})$$を求めてください。
多変数多項式の畳み込みは、１変数ずつ順に FFT する方法と、１変数多項式の畳み込み1回に帰着する方法の2つがあります。この記事では後者を使います。どの変数で切り捨てたか忘れることで計算量改善するので、前者では難しいと思います。前者は、多変数巡回畳み込み \(fg \bmod (x_1^{n_1}-1,\ldots,x_K^{n_K}-1)\) を本記事と同様の計算量で求める際に活躍します（演習問題２）。
\(fg\) は \(x_i\) について高々 \(2n_i\) 次式になりますから、\(fg\) を FFT で計算するには
$$\sum a_{i_1,i_2,\ldots,i_K}x^{i_1+2n_1i_2+(2n_1)(2n_2)i_3+\cdots +(2n_1)\cdots(2n_{K-1})i_K},$$
$$\sum b_{i_1,i_2,\ldots,i_K}x^{i_1+2n_1i_2+(2n_1)(2n_2)i_3+\cdots +(2n_1)\cdots(2n_{K-1})i_K}$$
を畳み込めばよいです。それぞれ、\(x\) の 高々\(2^{K-1}N\)次式（ただし\(N:=\prod n_i\)）ですから、計算量は
$$O( {\color{red} {2^K}} N \log (2^KN))$$ です。指数部分 \(O(2^K)\) を多項式に改善するのが目標です。
繰り上がりの回数を求める
\(O(2^K)\) を落とすために \(\sum a_{i_1,i_2,\ldots,i_K}x^{i_1+n_1i_2+\cdots +n_1 \cdots n_{K-1}i_K}, \sum b_{j_1,j_2,\ldots,j_K}x^{j_1+n_1j_2+\cdots +n_1 \cdots n_{K-1}j_K}\) として畳み込みたいのですが、繰り上がりが起きてしまいます。つまり、\(n_k \leq i_k+j_k < 2n_k\) のとき本来 \(x_k^{i_k+j_k}\) である項が \(x_{k+1}^{(i_k+j_k)/n_k}\) と区別できません。
そこで \(\chi(i+j)-\chi(i)-\chi(j)\) が繰り上がりの回数と等しくなるような関数 \(\chi\) を考えます。ただし、$$i = \sum_k i_k \prod_{l < k} n_{l} \quad (0 \leq i_k < n_k)$$
$$j = \sum_k j_k \prod_{l < k} n_{l} \quad (0 \leq j_k < n_k)$$
と置きました。
\begin{eqnarray}
\left\lfloor \frac{i_k + j_k}{n_k} \right\rfloor
=
\begin{cases}
1 & ( 繰り上がり発生 ) \\
0 & ( 繰り上がりしない)
\end{cases}
\end{eqnarray}
であることに着目します。\(\sum_{k < M}i_k \prod_{l < k} n_{l} < \prod_{1 \leq i \leq M} n_i\) なので、
\begin{eqnarray}
\left\lfloor \frac{i + j}{n_1n_2 \cdots n_k} \right\rfloor
-\left\lfloor \frac{i}{n_1n_2 \cdots n_k} \right\rfloor
-\left\lfloor \frac{j}{n_1n_2 \cdots n_k} \right\rfloor =
\begin{cases}
1 & ( i_k+j_k で繰り上がり発生 ) \\
0 & ( i_k+j_k で繰り上がりしない)
\end{cases}
\end{eqnarray}
となります。
つまり、
$$\chi(x)=\left\lfloor \frac{x}{n_1} \right\rfloor + \left\lfloor \frac{x}{n_1n_2} \right\rfloor + \cdots+ \left\lfloor \frac{x}{n_1n_2\cdots n_{K-1}} \right\rfloor $$が求めたい関数になっています。
$$F:=\sum a_{i_1,\ldots,i_K} t^{\chi(i_1+n_1i_2+\cdots+n_1\cdots n_{K-1} i_K)} x^{i_1+n_1i_2+\cdots+n_1\cdots n_{K-1} i_K},$$$$G:=\sum b_{j_1,\ldots,j_K} t^{\chi(i_1+n_1i_2+\cdots+n_1\cdots n_{K-1} i_K)} x^{i_1+n_1i_2+\cdots+n_1\cdots n_{K-1} i_K}$$とした 2 式を畳み込みます。高々 \(K\) 個の変数でしか繰り上がりが起きないので、\(\chi\) の畳み込み後の期待される値からの減少値は高々 \(K\) です。つまり、\(t^{K+1} = 1\) としても、繰り上がりを発見できます。どの変数で繰り上がりが起きたかは分からなくなってしまいましたが、その情報は不要です。
以上の計算は \(t, x\) を 1 変数に埋め込んで FFT することで \(O(K N\log(KN)) = O(KN \log(N))\) でできます。計算量の等号には \(N = N_1N_2\cdots N_K\) より、\(K = O(\log N)\) であることを使いました。
実装
\(t\) の冪指数ごとに分けて FFT することで求めることもできます。NTT のように DFT の長さに制約がある場合、こちらの方が長さを節約できて良いかもしれません。よすぽジャッジ[2]はギリギリこちらでしか通りません。\(F,G\) について、\(t\) の冪指数が \(0,1,\ldots,K\) の項を集めた多項式を \(t^0F_0,t^1F_1,\ldots,t^KF_K\), \(t^0G_0,t^1G_1,\ldots,t^KG_K\) とします。 \(F_0,\ldots,F_K,G_0,\ldots,F_K\) の FFT は \(O(KN\log(N))\) でできます。各 \(k\) の \(t^k\sum F_iG_{k-i}\) の FFT は畳み込み定理より合計 \(O(K^2N)\) 、IFFT は \(O(KN \log(N))\) です。すべて合計して \(O(KN\log(N)+K^2N) = O(KN \log(N))\) になりました。
参考文献
- rushcheyo的博客:集训队互测 2021 Round #1 题解
- 大変参考にしました。
- よすぽジャッジ：multivariate convolution
- コードを verify できます。私のコード。
Comments