趣旨
マシュマロとお題箱でイラストや記事のお題を募集しております。今回は「$p$進数上のFFTについて(原始根の存在や発見方法について)」という記事のお題をいただいたので、それに従って記事を書かせていただきます。
$p$進数の大まかなイメージと通常のFFT（高速フーリエ変換）とNTT（数論的変換）を前提知識として、対象は競技プログラミングに興味がある人を想定していきます。
前置き
$p$を素数とします。$p-1$の$2$進加法付値（$2$で割り切れる回数）を$d_p$と置きます。
原始根と言ったら普通は$p$進数体$\mathbb{Q}_p$の要素ではなく有限体$\mathbb{F}_p$などの要素を指しますが、質問の趣旨から察するに$\mathbb{Q}_p$における$1$の原始$p-1$乗根のことを指していると考えましょう。
FFTをするには$\mathbb{Q}_p$における$1$の原始$2^{d_p}$乗根さえ見つければ良いのですが、そのためには$\mathbb{Q}_p$における$1$の原始$p-1$乗根$z$を$1$つ探して$z^{(p-1)/2^{d_p}}$を考えれば良いわけです。
ではどうやって$\mathbb{Q}_p$における$1$の原始$p-1$乗根を探すのか、そもそも何故存在するのか、というのが今回のお題ですね。
なお実用上は$p$が$32$bit型整数に収まる範囲でありかつ$2^{d_p}$もまた十分大きいことが望ましいです。何故ならば、FFTによって直接計算できる畳み込みは位数$2^{d_p}$の巡回畳み込み（位数$2^{d_p}$の巡回群の群環における乗法）で、$d_p > 1$ならば巡回畳み込みを制限することで高々$2^{d_p-1}-1$次多項式の乗算も計算できるようになるからです。$d_p$が大きいほど、より次数の高い多項式の乗算が処理できるようになるわけですね。
そして愚直な多項式乗算で扱えないほど次数が高い状況こそ意味があるということを踏まえると、$d_p$は$16$以上であることが望ましいです。というわけで結果的に$p < 2^{32} \leq 2^{2d_p}$という関係式が成り立つことが望ましく、この$p < 2^{2d_p}$という関係式を満たす$p$をProth素数と言います。
そんなわけで$\mathbb{F}_p$上でNTTをしたり$\mathbb{Q}_p$上でFFTをしたりする上では実用上$p$がProth素数であってほしいわけですが、今回は一般論を説明するためそのような仮定をしないことに注意しておきます。
例題
そもそも何で$\mathbb{Q}_p$上でFFTなんて考えるのかが気になって先に進めなくなってしまう方は、例えば以下のような競技プログラミングの問題を考えてみましょう。
問題文
正整数$N$と$M$と非負整数係数$N$次多項式$F = \sum_{i=0}^{N} F_i x^i$が与えられます。素数$998244353$を$P$と置きます。$0 \leq n < P^2$を満たし$P$で割った余りが互いに異なる整数$n$を$M$個選び、それぞれの$n$に対し$F(n)$を$P^2 = 996491788296388609$で割った余りを求めてください。
制約
- $N$は$0 \leq N \leq 4 \times 10^5$を満たす整数である。
- $M$は$1 \leq M \leq 4 \times 10^5$を満たす整数である。
- $N$未満の任意の非負整数$i$に対し$F_i$は$0 \leq F_i < P^2$を満たす整数である。
- $F_N$は$0 < F_N < P^2$を満たす整数である。
実行時間制限 2 sec / メモリ制限 1024 MB
備考
$(n,F(n) \bmod P^2)$を$M$組出力する特殊ジャッジ問題（正解が複数ある問題）です。$n$を$P$で割った余りが異なるという制約を課しているのは、そうしないとFFTと関係なく簡単に解けるからです（考えてみてください）。
なお$P < 2^{32} < P^2 < 2^{64} < P^3$であるので、$0 \leq n,m < P^2$を満たす整数$n,m$に対して$64$bit整数を用いて愚直に法$P^2$乗算を行うと途中計算でオーバーフローする危険性があリます。そこでlong multiplication（位取り記法による乗法の筆算アルゴリズム）を用いたオーバーフロー回避テクニック \[ \begin{array}{l} nm \bmod P^2 = \\
\left( (n \bmod P)(m \bmod P) + \left( \left( \left( \left\lfloor \frac{n}{P} \right\rfloor (m \bmod P) \right) \bmod P \right) + \left( (n \bmod P) \left( \left\lfloor \frac{m}{P} \right\rfloor \right) \bmod P \right) \right) P \right) \bmod P^2 \end{array} \] で処理しましょう。
いわゆる任意mod FFT（中国剰余定理で有限個のProth素数に帰着させるFFT）と多点評価（中国剰余定理と分割統治FFTで複数点での代入を高速に計算するアルゴリズム）の組み合わせなら$O(N \log_2 N + M(\log_2 M)^2)$などで解くことが可能なので高速な言語なら間に合いそうですが、定数倍が重いので念のためより高速な解法を検討したいところです。
そこで実は$\mathbb{Q}_P$上のFFT（より正確には$P$進法の精度$2$桁で打ち切ったもの）を考えることでこの問題が$O(\max \{N,M\} \log_2 \max \{N,M\})$で解けるので、ぜひ考えてみてください。
もちろんFFTなんて使わずにもっと高速に解ける、という人もいるかもしれませんが許してください。
それでは例題はこの辺にしておいて、早速$p$進数のFFTについて説明していきます。
$1$の原始$p-1$乗根
環$R$と正整数$n$に対し、$R$における$1$の原始$n$乗根とは、$a^e = 1$を満たす最小の正整数$e$が$n$であるような要素$a \in R$のことです。特に$1$の原始$n$乗根は$1$の$n$乗根になります。
$\mathbb{F}_p$における$1$の原始$p-1$乗根は原始根とも呼ばれ、その存在は円分多項式などを用いて代数的に証明することが可能です。今回は$\mathbb{Q}_p$の話なので、ここは認めていきましょう。
一方で本題の$\mathbb{Q}_p$における$1$の原始$p-1$乗根はどうでしょうか？ 実はこちらも存在することが知られています。より強く、mod $p$を取る操作によって$\mathbb{Q}_p$における$1$の原始$p-1$乗根と$\mathbb{F}_p$の原始根が一対一に対応します（$\mathbb{Q}_p$における$1$の原始$p-1$乗根は$p$進整数環$\mathbb{Z}_p \subset \mathbb{Q}_p$に属すので、mod $p$を取る操作が定義されます）。
もっと言うと、Teichmuller embedding（またはTeichmuller lifting）と呼ばれる、mod $p$を取る操作と逆向きの単射群準同型（環準同型の制限ではないことに注意）
\[ \mathbb{F}_p^{\times} \hookrightarrow \mathbb{Z}_p^{\times} \]
が定義され、$\mathbb{F}_p^{\times}$の$p-1$個の要素（Fermatの小定理により全て$1$の$p-1$乗根）が$p-1$個の$1$の$p-1$乗根に送られ、その単射性から元々原始根だったものは行き先も$1$の原始$p-1$乗根になるというわけです。なお本当はTeichmullerのuは点々がつくのですが、表示の都合ご容赦ください。
ではこのTeichmuller embeddingの計算方法が分かればよいですね。そこで役立つのが次に説明するHenselの補題です。
$p$進Newton法
Henselの補題
$f$を$\mathbb{Z}_p$係数多項式とし、$\mathbb{F}_p$係数多項式$f \bmod p$が重根でない根$\zeta \in \mathbb{F}_p$を持つとし、$z_0 \in \mathbb{Z}_p$を$z_0 \bmod p = \zeta$を満たす要素とする。以下の漸化式で$\mathbb{Z}_p$の要素の列$(z_i)_{i=1}^{\infty}$を定める。 \[ z_i := z_{i-1} - \frac{f(z_{i-1})}{f’(z_{i-1})} \] この時、$p$進距離に関して$(z_i)_{i=1}^{\infty}$はただ１つの$p$進整数$z \in \mathbb{Z}_p$に収束し、以下が成り立つ：
(1) $z \bmod p = \zeta$である。
(2) 任意の$i \in \mathbb{N}$に対し、$z \bmod p^{2^i} = z_i \mod p^{2^i}$である。
(3) $f(z) = 0$である。
この漸化式による構成方法は実数に対するNewton法と同じなので、$p$進Newton法と呼ばれています。実数以外でもNewton法を使えることは、FFTで$\exp(X)$や$\frac{1}{1+X}$の合成を計算する方法をご存知なら納得しやすいですね。
さて$f(x) = x^{p-1} - 1$とすれば、$\mathbb{F}_p$の原始根$\zeta$が条件を満たします。重根にならないことは、実数の二次関数とかの時と同様に形式的に微分をすれば確認できますね（実数の性質は使っていません）。
というわけで、適当に固定した精度パラメータ$d \in \mathbb{N}$に対し$\mathbb{Q}_p$における$1$の原始$p-1$乗根$z$を$d$桁精度で（つまり$\mathbb{Z}/p^d \mathbb{Z}$の要素として厳密に）求めるには以下の手順を踏めば良いです。
- $\mathbb{F}_p$の原始根$\zeta$の代表元となる$z_0 \in \mathbb{Z}$を$0 < z_0 < p$の範囲で探す（全探索や乱択で高速に求まります）。
- $f(x) = x^{p-1} - 1$に対する先程の漸化式で$z_{\lceil \log_2 d \rceil}$までを法$p^d$で求める（分母に来るのは法$p^d$の可逆元であり、Eulerの定理＋繰り返し二乗法なりEuclidの互除法なり法$p$逆元前計算＋Taylor展開なりでいずれも高速に逆元計算できます）。
- $z_{\lceil \log_2 d \rceil}$を返す。
例えば整数係数多項式の$\mathbb{Q}_p$上でのFFTを法$p^2$で求めたければ$d = 2$とすれば良いです。$\lceil \log_2 d \rceil = 1$なので非常に高速ですね。
また$p$がコンパイル時定数で実行時に$\mathbb{Q}_p$上でFFTを処理したいならば、$1$の原始$p-1$乗根計算は実行時にせずとも前計算しておいてソースコードに埋め込めば良いです。となると実行時間には影響しなくなるので、競技プログラミングであれば提出コードを高速な言語で書くとしても前計算部分は書きやすい言語で書くという手があります。
実装
というわけでpythonで$\mathbb{Q}_p$における$1$の原始$p-1$乗根を$d$桁精度で求めてみましょう。
#素数pと正整数dに対しQ_pにおける1の原始p-1乗根を法p^dで求める。
def PrimitiveRoot(p,d):
assert PrimeFactor(p)==[p] and d>0
maximal_divisor=MaximalDivisor(p-1)
for z in range(1,p):
if all(pow(z,d,p)!=1 for d in maximal_divisor):
break
base=p**d
power=1
while power<d:
power<<=1
z=(z-(pow(z,p-1,base)-1)*pow((p-1)*pow(z,p-2,base),-1,base))%base
assert pow(z,p-1,base)==1 and all(pow(z,d,base)!=1 for d in maximal_divisor)
print(f"primitive root of 1 in Q_{p} modulo {p}^{d} = {base}: {z}")
return z
#正整数nの非自明な約数のうち極大なものを列挙する。
def MaximalDivisor(n):
answer=[n//p for p in PrimeFactor(n)]
print(f"maximal non-trivial divisors of {n}: {answer}")
return answer
#正整数nの素因数を重複なく列挙する。
def PrimeFactor(n):
answer=[]
p=2
while p*p<=n:
if n%p==0:
answer.append(p)
n//=p
while n%p==0:
n//=p
p+=1
if n>1:
answer.append(n)
return answer
print(PrimitiveRoot(998244353,2))
#出力:
#maximal non-trivial divisors of 998244352: [499122176, 142606336, 58720256]
#primitive root of 1 in Q_998244353 modulo 998244353^2 = 996491788296388609: 82941140536643239
例題の解説
$\max \{N,M\} \leq 4 \times 10^5 < 2^{19} < 2^{23} = 2^{d_P}$であることに注意すると、$F$の$\mathbb{Q}_P$上のFFTを法$P^2$で計算することで、$F$に例えば$1$の$2^{19}$乗根$2^{19}$個を代入した値を法$P^2$で求めることができます。
これにて確かに$O(\max \{N,M\} \log_2 \max \{N,M\})$で解くことができました。
FFTと多項式乗算を用語として混同していると思い付きにくく、FFTそのものが特定の評価点に対し分割統治を用いた多点評価アルゴリズムであることをきちんと理解しているかを問う問題でした。この見方だとIFFTは多項式補間アルゴリズムなわけですね。
特定の評価点に対する多点評価と多項式補間が同じアルゴリズムでできることは一見すると不思議ですが、有限巡回群のPontryagin双対が群同型であることを思い出すと自然に理解することが可能です。
自分用覚え書き
github pagesで数式を使う際は、閉じ括弧と半角アンダーバーの間に半角スラッシュを入れないとmarkdownからhtmlへの翻訳の仕様でバグるらしい。（波括弧でもたまにバグる）
数式環境をくくる大括弧と数式環境内の中括弧や半角空白や濃度のシャープは半角スラッシュ２つ、改行は半角スラッシュ５つ、その他は基本的に半角スラッシュ１つで良さそう。htmlとmathjaxの関係が難しい。
inline数式内の縦棒はlvertなどを使わないとmarkdownの表と認識される。
大括弧の数式環境はアンパサンドで縦位置を揃えられない。
arrayとalignはどちらも使える。arrayが使えない気がしてしまったのは中括弧につける半角スラッシュの個数を間違えていたから。alignはアンパサンドを等号の右側につけて半角空白を２つ入れると幅がちょうど良さそう。arrayなしだとアンパサンドは使えない。