FFT常数优化(共轭优化)

时间:2023-03-09 18:23:15
FFT常数优化(共轭优化)

最近闲着无聊研究了下\(FFT\)的常数优化,大概就是各种\(3\)次变\(2or1.5\)次之类的,不过没见过啥题卡这个的吧。

关于\(FFT\)可以看这里:浅谈FFT&NTT

关于复数

设\(x=a+bi\),其中\(i\)是虚数单位,那么我们用\(\bar x\)表示\(x\)的共轭复数,即\(\bar x=a-bi\)。

共轭复数有一个这样的性质:

\[\overline{ab}=\bar a \cdot \bar b
\]

证明展开就好了,这个是下面优化的关键。

设\(\omega_n\)为\(n\)阶单位根,则\(\overline{\omega _n^{x}}=\omega_{n}^{-x}\)。

idft变dft

设\(f(x)=\sum_{i=0}^{n-1}a_ix^i\),注意到:

\[n\cdot[j]{\rm idft}(f)=\sum_{i=0}^{n-1}\omega_{n}^{-ij}a_i=a_0+\sum_{i=1}^{n-1}\omega_{n}^{ij}a_{n-i}
\]

也就是说我们如果进行一次std::reverse(a+1,a+n),然后dft(a),在除以\(n\),我们就完成了一次\(idft\)。

多项式乘法优化 1

给出多项式\(f(x)=\sum_{i=0}^{n}a_ix^i,g(x)=\sum_{i=0}^{m}b_ix^i\),求其卷积。

这里最开始介绍一种非常简洁的优化方法,构造多项式\(h(x)\):

\[h(x)=f(x)+ig(x)
\]

\[h^2(x)=f^2(x)-g^2(x)+2if(x)g(x)
\]

那么我们只需要取\(h^2(x)\)的虚部除以\(2\)就是答案,这只需要做两次\(FFT\)。

多项式乘法优化 2

这个和上面的关联不大,设\(X_i\)表示多项式\(F(x)\)\(dft\)之后的系数,\(a_i\)表示\(dft\)之前的系数,设\(F(x)\)为\(n\)项的多项式,且\(n=2^k\),注意到:

\[X_i=\sum_{j=0}^{n-1}a_j\omega_{n}^{ij},X_{n-i}=\sum_{j=0}^{n-1}a_j\omega_{n}^{-ij}
\]

即:\(X_i=\overline{X_{n-i}}\)。

这实质上是因为\(F\)没有虚部的原因,我们换一个有虚部的多项式试试:

\[X_i=\sum_{j=0}^{n-1}(a_j+ib_j)\omega_{n}^{ij}\\
X_{n-i}=\sum_{j=0}^{n-1}(a_j+ib_j)\omega_{n}^{-ij}\\
\overline{X_{n-i}}=\sum_{j=0}^{n-1}(a_j-ib_j)\omega_{n}^{ij}\\
\]

等等,我们发现第一个式子和第三个式子很像,两式相加减可以得到:

\[X_i+\overline{X_{n-i}}=2\sum_{j=0}^{n-1}a_j\omega_{n}^{ij}\\
X_i-\overline{X_{n-i}}=2i\sum_{j=0}^{n-1}b_j\omega_{n}^{ij}
\]

注意到等式右边就是\(a\) \(dft\)完之后的结果,那么对于多项式\(F(x),G(x)\),我们可以构造一个函数然后\(dft\)一次,然后\(O(n)\)得到两个多项式\(dft\)之后的结果,总共只用了一次\(FFT\)。


当然这个玩意也可以这样用:假设我们现在想求\(dft(F(x))\),我们把\(F(x)\)奇偶分类,构造多项式:

\[g(x)=\sum_{i=0}^{n/2-1}(a_{2i}+ia_{2i+1})x^i
\]

然后相当于是\(0.5\)次\(FFT\)来完成这个事,设\(dft(g(x))\)每一项为\(X_i\),\(dft(F(x))\)每一项为\(Y_i\),那么推一下可以得到:

\[Y_i=\frac{X_i+\overline{X_{n/2-i}}}{2}-2\omega_{n}^i(X_i-\overline{X_{n/2-i}})
\]

注意这里只有\(i\in [0,n/2)\)的值,\(Y_{n/2}​\)特殊处理一下,后面的可以通过前面得到。

MTT常数优化

\(\rm MTT\)就是拆系数\(\rm FFT\),设多项式\(s(x),t(x)\),我们要算\(s(x)t(x)\),模数任意。

我们拆系数,设拆完了之后是\(s(x)=a(x)+b(x)\cdot p,t(x)=c(x)+d(x)\cdot p\)。

构造\(F(x)=a(x)+i\cdot b(x)\),\(G(x)=c(x)+i\cdot d(x)\)。

那么有:

\[\begin{align}
&F(\omega_n^j)=\sum_{i=0}^{n-1}(a_i+ib_i)\omega_n^{ij}\\
&F(\omega_n^{-j})=\sum_{i=0}^{n-1}(a_i+ib_i)\omega_n^{-ij}\\
&\overline{F(\omega_n^{-j})}=\sum_{i=0}^{n-1}(a_i-ib_i)\omega_n^{ij}\\
\end{align}
\]

那么相加减可得\(a(x),b(x)\)的\(dft\)。

令\(h(x)={\rm dft}(a(x))\cdot {\rm dft}(G(x))={\rm dft}(a(x)\cdot G(x))={\rm dft}(a(x)c(x)+i\cdot a(x)d(x))\)。

那么我们\(idft\)一次\(h(x)\)就可以得到\(a(x)c(x),a(x)d(x)\)。

同理可以得到\(b(x)c(x),b(x)d(x)\),一共\(4\)次\(dft\)。

代码长这样:

void mul(int *r,int *s,int *t,int len) {
for(N=1,bit=0;N<len;N<<=1,bit++);
for(int i=1;i<N;i++) pos[i]=pos[i>>1]>>1|((i&1)<<(bit-1));
for(int i=0;i<N;i++) g[0][i]=cp(r[i]&all,r[i]>>15),g[1][i]=cp(s[i]&all,s[i]>>15);
fft(g[0]),fft(g[1]);
for(int i=0;i<N;i++) {
int j=(N-i)&(N-1);
g[2][j]=(g[0][i]+conj(g[0][j]))*cp(0.5,0)*g[1][i];
g[3][j]=(g[0][i]-conj(g[0][j]))*cp(0,-0.5)*g[1][i];
}fft(g[2]),fft(g[3]);
for(int i=0;i<N;i++) g[2][i]=g[2][i]/N,g[3][i]=g[3][i]/N;
for(int i=0;i<N;i++) {
ll pp=g[2][i].r+0.5,x=g[2][i].i+0.5,y=g[3][i].r+0.5,z=g[3][i].i+0.5;
t[i]=(pp%p+(((x+y)%p)<<15)+((z%p)<<30))%p;
}
}