Skip to content →

多変数畳み込み(切り捨て)のアルゴリズム

中国のブログ[1]に面白いアルゴリズムがあったので、紹介します。

次のような多変数畳み込み(切り捨て)を計算するアルゴリズムを構築します:

\(K\) 変数の多項式
$$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})$$を求めてください。

多変数多項式の畳み込みは、1変数ずつ順に FFT する方法と、1変数多項式の畳み込み1回に帰着する方法の2つがあります。この記事では後者を使います(どの変数で切り捨てたか忘れることで計算量改善するので、前者では難しいと思います)。

\(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}\) と区別できません。

そこで繰り上がりを管理するための変数 \(t\) と重みづけの関数
$$\chi(n)=\left\lfloor \frac{n}{n_1} \right\rfloor + \left\lfloor \frac{n}{n_1n_2} \right\rfloor + \cdots+ \left\lfloor \frac{n}{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 式を畳み込みます。このアルゴリズムの肝は関数 \(\chi\) です。\(\left\lfloor \frac{n}{n_1n_2\cdots n_{k}} \right\rfloor \) が変数 \(x_k\) で繰り上がりが起きたかどうか判別するための項です。繰り上がりが起きたとき、つまり \(n_k \leq i_k + j_k < 2n_k\) のとき、この項だけが畳み込み後に期待される値から \(1\) 減少して$$\left\lfloor \frac{n_1\cdots n_{k-1}(i_k+j_k)}{n_1\cdots n_{k}}\right\rfloor – \left(\left\lfloor \frac{n_1\cdots n_{k-1}i_k}{n_1\cdots n_{k}}\right\rfloor + \left\lfloor \frac{n_1\cdots n_{k-1}j_k}{n_1\cdots n_{k}}\right\rfloor\right) = -1$$となり、さもなくば、どの項も変化なく、特に$$\left\lfloor \frac{n_1\cdots n_{k-1}(i_k+j_k)}{n_1\cdots n_{k}}\right\rfloor – \left(\left\lfloor \frac{n_1\cdots n_{k-1}i_k}{n_1\cdots n_{k}}\right\rfloor + \left\lfloor \frac{n_1\cdots n_{k-1}j_k}{n_1\cdots n_{k}}\right\rfloor\right) = 0$$です。高々 \(K\) 個の変数でしか繰り上がりが起きないので、\(\chi\) の畳み込み後の期待される値からの減少値は高々 \(K\) です。つまり、\(t^{K+1} = 1\) としても、繰り上がりを発見できます。どの変数で繰り上がりが起きたかは分からなくなってしまいましたが、だからこそ計算量削減が期待できます。\(\chi\) の関数の形は突拍子がないように感じましたが、(1) 繰り上がりが起きた時だけ、減少する (2) 減少幅は 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))\) になりました。

参考文献

  1. rushcheyo的博客:集训队互测 2021 Round #1 题解
    • 大変参考にしました。
  2. よすぽジャッジ:multivariate convolution
多変数の畳み込みを1変数の畳み込みに直すのは、単なる実装の工夫だと思っていました。計算量改善に役立つとは!

Published in アルゴリズム

Comments

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

CAPTCHA