模意义下的FFT——快速数论变换（Number-Theoretic Transform，NTT）学习

2019-04-13 15:31发布

站内文章 / 模拟电子
12757 0 1164

前置技能

前言

  • 快速数论变换（Number-Theoretic Transform，NTT），实际上就是模意义下的FFT。
  • 你可能会问：既然已经有FFT了，那要NTT有何用？
  • 实际上，FFT的弊病很多。
  • 最致命的弱点在于：FFT涉及复数运算，而当数字一大，难免出现精度问题。
  • 况且，我们在计算主N次单位根ωnomega_n时，还要用到cos⁡cos、sin⁡sin这类超越函数，而其时间复杂度远远大于普通的加法、乘法。
  • 再者，随着数的增大，答案可能也会变得超级大，题目会让我们输出答案对某个质数取模的结果。而FFT的复数是由两个实数构成（分别表示实部和虚部），实数类型不好取模；而算出最终答案再取模，又可能会爆long long。
  • 如果我们能用整数替换掉N次单位复数根，事情就变得非常好看了。
  • 我们来回忆一下FFT能在O(nlog2n)O(nlog_2n)的时间内变换用到了N次单位根的哪些性质：
  • ωn0=ωnn=1omega_n^0=omega_n^n=1
  • ωn0,ωn1,...,ωnn−1omega_n^0,omega_n^1,...,omega_n^{n-1}互不相同（这样就可以保证没有重复的点值对），并且为一个公比为ωn1omega_n^1的等比数列。
  • 消去引理：ωdndi=ωniomega_{dn}^{di}=omega_n^i
  • 正因为单位根ωomega满足这些性质，我们才能舒服地推出折半引理、求和引理，进而完美地证明后面采取的每一种策略都是正确的。

原根

  • 现在我们要在数论中寻找满足这三个性质的数，首先来介绍原根的概念，根据费马小定理我们知道，对于一个素数p，有下面这样的关系：ap−1≡1 (mod p)a^{p-1}equiv1 (mod p)
  • 这和主N次单位根十分相似。而 p 的原根 g 定义为使得∀0≤i<j<p−1,gi̸≡gj (mod p)forall 0≤i<j<p-1,g^i otequiv g^j (mod p)的数。
  • 我们取素数p=k⋅2x+1p=k·2^x+1，并找到它的原根g。譬如素数998244353=119∗223+1998244353=119*2^{23}+1，3就是它的一个原根。
  • 令gn≡g(p−1)/n (mod p)g_nequiv g^{(p-1)/n} (mod p)，即原根的公比。囿于我们人为地将n扩大到2的某个次幂，所以只要p=k⋅2x+1p=k·2^x+1中的x够大，(p-1)/n就可以为整数。
  • 这样，我们就可以使得gn0,gn1,...,gnn−1 (mod p)g_n^0,g_n^1,...,g_n^{n-1} (mod p)互不相同，并且gn0≡gnn≡1 (mod p)g_n^0equiv g_n^nequiv1 (mod p)。而gdndi≡gdi(p−1)/dn≡gi(p−1)/n≡gnig_{dn}^{di}equiv g^{di(p-1)/dn}equiv g^{i(p-1)/n}equiv g_n^i，显然也满足消去引理。而折半引理、求和引理同样也可以手玩证得。

具体实现

  • 只需要将ωniomega_n^i替换成gnig_n^i即可。
  • 我们在做IDFT的时i，也是把ωn−iomega_n^{-i}替换成gn−ig_n^{-i}，也就是gnn−ig_n^{n-i}，当然其实就是gnig_n^i的相反数；然后和FFT的IDFT一样，最后再给答案除以n。
  • 其他的跟FFT相比，没什么变化。
  • 具体参考Code。

Code

#define go(i,a,b) for(i=a;i< b;i++) ll fpow(ll x,ll y)//快速幂 { ll ans=1; for(;y;y>>=1) { if(y&1) T(ans,x); T(x,x); } return ans; } ll inv(ll x) {return fpow(x,mo-2);}//求x对mo的逆元 void NTT(ll *a,bool flag)//flag=0表示DFT，flag=1表示IDFT { go(i,0,m) if(i>tf[i])swap(a[i],a[tf[i]]);//蝶形变换 ll u; for(int n=2;n<=m;n<<=1) { int h=n>>1; ll w_n=fpow(pr,(mo-1)/n), w=1;//pr表示原根；w_n表示gn，即数列g的公比 if(flag) w_n=inv(w_n); go(i,0,h) { for(j=i;j<m;j+=n) { u=a[j+h]*w%mo; a[j+h]=(a[j]-u+mo)%mo; a[j]=(a[j]+u)%mo; } w=w*w_n%mo; } } if(flag) go(i,0,m) T(a[i],inv(m));//如果是IDFT，最后答案再除m }

任意模数NTT

  • 前面说了，NTT需要的模数是p=k⋅2x+1p=k·2^x+1形式的素数，但是在实际生活中，要求的模数可能不是这样的形式，甚至是一个合数！
  • 那怎么破呢？考虑选取k个可以用上面形式表示的素数，一般是选取这3个：998244353 , 1004535809 , 469762049998244353 , 1004535809 , 469762049或 9985661441=235∗222+1 9985661441=235*2^{22}+1。
  • 我们发现3均为它们的原根，那么对每个模数分别做一波NTT，最后使用中国剩余定理合并。
  • 不妨设我们选取的k个模数为p1,p2,...,pkp_1,p_2,...,p_k

热门文章