はじめに
競技プログラミングで大きい整数の掛け算のアルゴリズムとしてFFTやNTTを勉強すると、これらのアルゴリズムは桁数$n$が大きくなっても使えるのかという疑問を抱きませんか？
複素数体を用いたFFTは$n$が大きくなっても浮動小数点数の精度を十分に大きくすれば正確に掛け算できると思いますが、どれだけ精度を確保すればよいか分からないので正確に計算するのに必要な計算量は分かりませんでした。
また競技プログラミングなどで出てくるNTTは、$p = 998244353$として$\mathbb{F}_p$を用いれば$n$が$10^6$くらいまでは正しく計算できますが、$n$が大きくなると都合の良い素数をどのようにして見つけてくればよいかかが分かりませんでした。
そこで私は$n$がどれだけ大きくなっても使えるアルゴリズムを探しました。するとSchönhage-Strassen (ショーンハーゲ・ストラッセン)法を見つけました。このアルゴリズムで$n$桁の整数同士の掛け算を$O(n \log n \log (\log n))$で計算できることを知ったので本記事で紹介したいと思います。
高速フーリエ変換
ショーンハーゲ・ストラッセン法でも高速フーリエ変換を使います。競技プログラミングで出てくる高速フーリエ変換は複素数体や素数体などの体を対象にしていましたが、ショーンハーゲ・ストラッセン法では可換環を対象にします。まず可換環での離散フーリエ変換を見ていきましょう。
可換環での離散フーリエ変換
以下を満たす可換環$R$と自然数$m$と$1$の原始$m$乗根$r \in R$を考えます。
- $mx = 1$となる$x \in R$が存在する。 (この$x$を$m^{-1}$とする)
- $\sum_{l=0}^{m-1}r^{kl} = 0$が$k=1,2,\ldots,m-1$で成り立つ。
これらの条件を満たす$(R, m, r)$に対して$(a_0,a_1,\ldots,a_{m-1}) \in R^m$を
$$b_k = \sum_{l=0}^{m-1}a_lr^{kl} \ (k=0,1,\ldots,m-1)$$
で与えられる$(b_0,b_1,\ldots,b_{m-1}) \in R^m$にする変換を離散フーリエ変換と呼びます。また$(b_0,b_1,\ldots,b_{m-1}) \in R^m$を
$$a_k = m^{-1}\sum_{l=0}^{m-1}b_lr^{(m-k)l} \ (k=0,1,\ldots,m-1)$$
で与えられる$(a_0,a_1,\ldots,a_{m-1}) \in R^m$にする変換を逆離散フーリエ変換と呼びます。離散フーリエ変換と逆離散フーリエ変換に対応する行列はそれぞれ
\begin{pmatrix}
r^{0 \cdot 0} & r^{0 \cdot 1} & \cdots & r^{0 \cdot (m-1)}\\
r^{1 \cdot 0} & r^{1 \cdot 1} & \cdots & r^{1 \cdot (m-1)}\\
\vdots & \vdots & \ddots & \vdots \\
r^{(m-1) \cdot 0} & r^{(m-1) \cdot 1} & \cdots & r^{(m-1) \cdot (m-1)}\\
\end{pmatrix}, m^{-1}\begin{pmatrix}
r^{m \cdot 0} & r^{m \cdot 1} & \cdots & r^{m \cdot (m-1)}\\
r^{(m-1) \cdot 0} & r^{(m-1) \cdot 1} & \cdots & r^{(m-1) \cdot (m-1)}\\
\vdots & \vdots & \ddots & \vdots \\
r^{1 \cdot 0} & r^{1 \cdot 1} & \cdots & r^{1 \cdot (m-1)}\\
\end{pmatrix}
です。2つ目の条件式から2つの積が$I_m$で逆離散フーリエ変換が離散フーリエ変換の逆変換になっていることが分かります。
ショーンハーゲ・ストラッセン法で用いる可換環
ショーンハーゲ・ストラッセン法では$(R,m,r)=(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},n,4^k)$に対する離散フーリエ変換を用います。ただし$k$は自然数で$n$は$2$の冪乗の自然数です。まず$4^k$が$1$の原始$n$乗根なことを確かめます。まず
$$4^{kn} = (2^{kn}+1)(2^{kn}-1)+1\equiv 1\ (mod \ 2^{kn} +1)$$
で$4^k$が$1$の$n$乗根なことが分かります。$n$が$2$の冪乗なことと
$$4^{k\frac{n}{2}} = 2^{kn}\equiv -1\ (mod \ 2^{kn} +1)$$
で$4^k$が$1$の$\frac{n}{2}$乗根でないことから、$4^k$が$1$の原始$n$乗根なことが分かります。
次に上の節であげた$2$つの条件を満たすことを確認します。
1つ目の条件は$n$が$2$の冪乗で$2^{kn}+1$が奇数で2つは互いに素なので成り立ちます。
2つ目の条件ですがまず$j = 1,2,\cdots,n-1$で$1 - 4^{kj}$が$2^{kn}+1$と互いに素なことを確認します。2つが$1$でない公約数$g$を持つと仮定します。$g$が$2$出ないことに注意します。以下の式が
$$4^{kj}\equiv 1,\ 2^{kn}\equiv -1\ (mod \ g) $$
成り立ちます。後者から上と同様の議論で$4^k$が$\mathbb{Z}/g\mathbb{Z}$で$1$の原始$n$乗根なことが分かります。よって前者から$j$は$n$の倍数となりますが$j = 1,2,\cdots,n-1$に矛盾します。$1 - 4^{kj}$が$2^{kn}+1$と互いに素なことが分かりました。
だから$\mathbb{Z}/(2^{kn}+1)\mathbb{Z}$で$(1 - 4^{kj})^{-1}$が存在します。ところで
$$(1 - 4^{kj})\sum_{l=0}^{n-1}4^{kjl} = 1 - 4^{kjn} = 0$$
が成り立ちます。両辺に$(1 - 4^{kj})^{-1}$をかけると
$$\sum_{l=0}^{n-1}4^{kjl} = 0$$
が言えます。
$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},n,4^k)$が2つの条件を満たすことを確認できました。
高速化
自然数$k$と$2$の冪乗の自然数$n$に対し$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},n,4^k)$の離散フーリエ変換の計算の高速化をみていきましょう。$(a_0,a_1,\ldots,a_{n-1}) \in (\mathbb{Z}/(2^{kn}+1)\mathbb{Z})^n$に対して
$$b_j = \sum_{l=0}^{n-1}a_l4^{kjl} \ (j=0,1,\ldots,n-1)$$
を計算します。
$$b_j = \sum_{l=0}^{n/2-1}a_{2l}4^{2kjl} +4^{kj}\sum_{l=0}^{n/2-1}a_{2l+1}4^{2kjl}$$
が成り立ちます。$\sum_{l=0}^{n/2-1}a_{2l}4^{2kjl}$と$\sum_{l=0}^{n/2-1}a_{2l+1}4^{2kjl}$はそれぞれ$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},\frac{n}{2},4^{2k})$の離散フーリエ変換を$(a_0,a_2,\ldots,a_{n-2})$と$(a_1,a_3,\ldots,a_{n-1})$に適用したもの第$j-\lfloor\frac{2j}{n}\rfloor\frac{n}{2}$成分であることに注意します。2つの離散フーリエ変換が求まっていると$b_j$はそれぞれ$2$回の$\mathbb{Z}/(2^{kn}+1)\mathbb{Z}$の演算で求めることができます。
よって$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},n,4^k)$の離散フーリエ変換は全体で$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},\frac{n}{2},4^{2k})$の離散フーリエ変換2回と$\mathbb{Z}/(2^{kn}+1)\mathbb{Z}$の演算$2n$回で求めることができます。
よって$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},\frac{n}{2^l},4^{2^lk})$の離散フーリエ変換を計算するのに必要な演算数を$l=0,1,\cdots, \log_2n$で$opr(\log_2n-l)$とすると
$$opr(l) = 2opr(l-1) + 2\cdot 2^{l} \ (l=1,2,\cdots, \log_2n)$$
が分かります。$opr(0)=0$と合わせて$opr(l) = l2^{l+1}$が分かります。
特に$l=\log_2n$のときを考え$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},n,4^k)$の離散フーリエ変換は$2n\log_2n$回の$\mathbb{Z}/(2^{kn}+1)\mathbb{Z}$の演算で計算できることが分かりました。
逆離散フーリエ変換は最後に$m^{-1}$をかける操作が加わりますがそれ以外は離散フーリエ変換と同じです。
アルゴリズムの内容
この章ではショーンハーゲ・ストラッセン法の内容を見ていきます。ショーンハーゲ・ストラッセン法でも整数の掛け算を多項式の掛け算に帰着し、多項式の掛け算を離散フーリエ変換を使って高速化します。整数の入出力は両方2進数とします。
多項式の掛け算への帰着
ショーンハーゲ・ストラッセン法では適切な$d$を選び整数を$2^d$進法で考えます。整数の$2^d$進法表示は$2$進法表示を$d$個ずつに分けると計算できます。
2つの$2^d$進数で$n$桁の整数の$2^d$進数表示をそれぞれ$a_{n-1}a_{n-2}\cdots a_0$、$b_{n-1}b_{n-2}\cdots b_0$とします。多項式$f$,$g$を
$$f(x) = \sum_{j=0}^{n-1}a_jx^j, g(x) = \sum_{j=0}^{n-1}b_jx^j$$
とすると2つの整数はそれぞれ$f(2^d)$,$g(2^d)$です。だから$f(2^d)g(2^d)$の2進数表示を計算すればよいです。これは$f$と$g$の積を計算した後適切に繰り上がりを処理すれば計算できます。
離散フーリエ変換を用いた多項式の掛け算
離散フーリエ変換は$m$次以下の多項式の係数を受け取って、$r^0,r^1,\cdots,r^{m-1}$を代入したものを計算していると見ることができます。また逆離散フーリエ変換は$m$次以下の多項式に$r^0,r^1,\cdots,r^{m-1}$を代入したものを受け取って、多項式の係数を計算していると見ることができます。
これを踏まえて、$m$次の多項式$f$と$g$の積は以下のように計算できます。
- $f (r^0), f (r^1),\cdots,f (r^{2m-1})$,$g (r^0), g (r^1),\cdots,g (r^{2m-1})$を離散フーリエ変換で計算
- $f (r^0)g(r^0), f (r^1)g(r^1),\cdots,f (r^{2m-1})g(r^{2m-1})$を計算
- 逆離散フーリエ変換をで$f$と$g$の積の係数を計算
ショーンハーゲ・ストラッセン法での適用
$2^d$進法で$\frac{n}{2}$桁の整数同士の掛け算を見ていきましょう。
2つの整数の$2^d$進数表示をそれぞれ$a_{\frac{n}{2}-1}a_{\frac{n}{2}-2}\cdots a_0$、$b_{\frac{n}{2}-1}b_{\frac{n}{2}-2}\cdots b_0$とします。$0\leqq a_j,b_j < 2^d \ (j = 0,1,\cdots,\frac{n}{2}-1)$です。
$(\mathbb{Z}/(2^{kn}+1)\mathbb{Z},n,4^k)$の$k$をどうするか考えます。
$\sum_{j=0}^{\frac{n}{2}-1}a_jx^j$と$\sum_{j=0}^{\frac{n}{2}-1}b_jx^j$の積を正しく計算するには積の係数が$2^{kn}+1$を超えない必要があります。係数は$n4^d$未満なので
$$2^{kn}+1 > n4^d$$
であれば十分です。これを満たす$k$に対し上の方法で$\mathbb{Z}/(2^{kn}+1)\mathbb{Z}$での多項式の積の係数を求めればよいです。
計算量
まず$\mathbb{Z}/(2^{kn}+1)\mathbb{Z}$での演算にかかる計算量を見ていきます。
足し算についてみていきます。整数の足し算にかかる計算量は$O(kn)$です。整数の足し算のあと$2^{kn}+1$で割らないといけませんが、和が$2^{kn}+1$を超えているときに$2^{kn}+1$を引けばよいので$O(kn)$です。
掛け算についてみていきます。整数の掛け算は一般の場合はショーンハーゲ・ストラッセン法を再帰的に使うので後で考えます。どちらかの数が$2$の冪乗の場合はもう片方の整数をシフトするだけで計算できるので$O(kn)$です。整数の掛け算のあと$2^{kn}+1$で割りますが、両方の整数が$2^{kn}$でない場合は積の桁数は$2kn$以下です。この場合、下$kn$桁から上$kn$桁を引いて負になったら$2^{kn}+1$を足したのが余りです。これは$O(kn)$で計算できます。両方の整数が$2^{kn}$の場合余りは$1$なので場合分けすればよいです。
ではショーンハーゲ・ストラッセン法の計算量を見ていきましょう。$2^d$進法で$n$桁の整数同士の掛け算を考えます。$n\to 2^{\lceil \log_2 n\rceil}$とすれば$n$が2の冪乗のときのみを考えればよいことが分かります。
まず$(\mathbb{Z}/(2^{2kn}+1)\mathbb{Z},2n,4^k)$での離散フーリエ変換の計算量を考えます。上で見た通り$4n(\log_2 n + 1)$回の$\mathbb{Z}/(2^{2kn}+1)\mathbb{Z}$の演算で計算できます。離散フーリエ変換の演算には足し算も掛け算も使いますが掛け算は片方が2の冪乗です。そのため全体で$O(kn^2\log n)$で計算できます。
次に離散フーリエ変換の出力を掛け合わせるステップの計算量を考えます。掛け算は$2n$回行わないといけないので桁数$kn$の整数同士の掛け算$2n$回及び余りを計算する$O(kn^2)$が入ります。
次に逆離散フーリエ変換ですが$(2n)^{-1}$を掛けるステップ以外は離散フーリエ変換と同じで計算量は$O(kn^2\log n)$です。$n$が2の冪乗なことに注意すると$(2n)^{-1} = 2^{4kn-1-\log_2 n}$です。そのため$(2n)^{-1}$を掛けるのは$2n$個合わせて$O(kn^2)$です。
最後に繰り上がりの処理ですが$O(kn^2)$でできます。
合わせて$O(kn^2\log n)$と桁数$kn$の整数同士の掛け算$2n$回です。
今まで$2^d$進数で考えていたので$2$進数に戻して考えます。$2$進数で$n$桁の自然数同士の掛け算を計算するのに必要な計算量を$T(n)$とすると
$$T(n) = O\left(k\left(\frac{n}{d}\right)^2\log \frac{n}{d}\right) + T\left(\frac{kn}{d}\right)O\left(\frac{n}{d}\right)$$
で計算できることがわかります。
$k$は$2^{\frac{kn}{2d}}+1 > \frac{n}{2d}4^d$を満たしていればよいので$kn = O(d(d + \log n))$にできます。$d$を$\log n$より大きくとるようにすると
$$T(n) = O\left(n\log \frac{n}{d}\right) + T(O(d))O\left(\frac{n}{d}\right)$$
です。$A$,$B$,$C$を自然対数より大きい正の定数として
$$T(n) = An\log \frac{n}{d} + C \cdot \frac{n}{d} T(Bd)$$
となります。
以下では$d$や再帰の仕方をどうするか考えます。展開を繰り返すと
$$T(n) = An\log \frac{n}{d_1} + An BC\log \frac{Bd_1}{d_2} + \cdots + AnB^lC^l\log \frac{Bd_l}{d_{l+1}} + B^lC^{l+1}\frac{n}{d_{l+1}}T(Bd_{l+1})$$
となります。$d_i = n^{ \frac{1}{B^{i}C^{i}}}$とすると
$$T(n) = A\left(1 - \frac{1}{BC}\right)(l+1)n\log n + An\log B \cdot BC\cdot \frac{B^lC^l - 1}{BC - 1} + B^lC^{l+1}n^{1 - \frac{1}{B^{l+1}C^{l+1}}}T\left(Bn^{\frac{1}{B^{l+1}C^{l+1}}}\right)$$
です。$d > \log n$に対応する条件は$\log (Bd_i) < d_{i+1}$で
$$\log B + \frac{\log n}{B^iC^i} < n^{\frac{1}{B^{i+1}C^{i+1}}}$$
となります。ところで
$$\log B + BC \log x = x$$
は$x > BC$に唯一の実数解 $\alpha$を持ちます。
$$n^{\frac{1}{B^{i+1}C^{i+1}}} > \alpha$$
であれば十分で$l$はこれを満たす最大の$i$とします。すると$l=O(\log (\log n))$で$Bn^{\frac{1}{B^{l+1}C^{l+1}}} = O(1)$になります。$Bn^{\frac{1}{B^{l+1}C^{l+1}}}$桁の掛け算を愚直に計算すると
$$T(n) = O(n \log n \log (\log n))$$
となります。以上より$O(n \log n \log (\log n))$の計算量で計算できます。
NTTとの比較
競技プログラミングででてくる$998244353$などの性質のよい素数を使って計算するいわゆるNTTとショーンハーゲ・ストラッセン法を比較してみましょう。
NTTで使っている体は例えば$F_{998244353}$で$15311432$が$1$の$2^{23}$乗根となっています。$998244353=119\cdot 2^{23}+1$なので$2^{23}$は$998244353$の約$0.01$倍です。
一方ショーンハーゲ・ストラッセン法で使っている可換環の$\mathbb{Z}/(2^{kn}+1)\mathbb{Z}$で$4^k$は$1$の$n$乗根となっています。$n \ll 2^{kn}+1$なことを考えるとNTTの方が性質のよい可換環を使っているように思えます。
では任意の$2$の冪乗の自然数$n$に対し$O(n)$の素数$p$と$F_p$での$1$の原始$n$乗根を持ってこれたとしましょう。これらを使ったNTTでは計算量が$O(n\log n (F_p\mbox{の演算にかかる計算量}))$となります。$F_p$での演算にかかる計算量は明らかに$O(\log n)$以上なのでショーンハーゲ・ストラッセン法よりは遅いことが分かります。
$n$が十分大きいときはブロック分けをして計算した方が効率的でそうなると上で見たような可換環の性質のよさは重要ではないのかなと思いました。また原始根が$2$の冪乗で掛け算がしやすいのもショーンハーゲ・ストラッセン法が高速な理由になっていそうです。
最後に
今回は桁数がどれだけ大きくても使える掛け算のアルゴリズムのショーンハーゲ・ストラッセン法をみました。長い間どういうアルゴリズムなのか気になっていたので今回勉強できてよかったです。
最近$O(n \log n)$で掛け算を計算できるアルゴリズムが見つかったそうです。難しそうですがいつか理解できたらと思います。
最後に宣伝になりますが計算の工夫を用い20桁くらいの大きい桁数の数同士の掛け算や割り算を人力で行う方法の記事も書いています。よかったらそちらも見てください。