【详解】如何写出 448b 的多项式乘法

· · 算法·理论

with Glacia_Official&le0n.

顺便记录一下 NTT 是怎么做的。前置知识:复数,单位根,矩阵,转置原理,原根。

DFT 干的是这样的事:给定多项式 F(x)=f_0+f_1x+\cdots+f_{2^m-1}x^{2^m-1},求出所有 F(\omega_{2^m}^i),IDFT 则是这个过程的逆。设 F(x)=G(x^2)+xH(x^2),其中 G,H 分别是 \deg<2^{m-1} 的多项式,由 F 分离偶数与奇数项系数得到。那么有 F(\omega_{2^m}^i)=G(\omega_{2^{m-1}}^{i\bmod 2^{m-1}})+\omega_{2^m}^i\cdot H(\omega_{2^{m-1}}^{i\bmod 2^{m-1}}),因此递归可以做到 O(2^mm)。改为非递归形式,如果将分治过程写出来就是

[0, 1, 2, 3, 4, 5, 6, 7]
[0, 2, 4, 6][1, 3, 5, 7]
[0, 4][2, 6][1, 5][3, 7]
[0][4][2][6][1][5][3][7]

那么最开始将 a_i 放在 a_{\text{rev}(i)} 的位置(其中 \text{rev} 代表二进制下反转第 0\sim m-1 位)。这个操作被称为位反转变换。接下来自底向上地合并信息即可完成 DFT。接下来考虑 IDFT 怎么做,自然可以自上而下地解,但是有一个更巧妙的方法。DFT 实际上是一个线性变换,把矩阵写出来是:

\begin{bmatrix} F(\omega_{2^m}^0)\\ F(\omega_{2^m}^1)\\ \vdots\\ F(\omega_{2^m}^{2^m-1}) \end{bmatrix}= \begin{bmatrix} \omega_{2^m}^0&\omega_{2^m}^0&\cdots&\omega_{2^m}^0\\ \omega_{2^m}^0&\omega_{2^m}^1&\cdots&\omega_{2^m}^{2^m-1}\\ \vdots&\vdots&\ddots&\vdots\\ \omega_{2^m}^0&\omega_{2^m}^{2^m-1}&\cdots&\omega_{2^m}^{(2^m-1)(2^m-1)} \end{bmatrix} \begin{bmatrix} f_0\\ f_1\\ \vdots\\ f_{2^m-1} \end{bmatrix}

考虑这个转移矩阵 \text{DFT} 的平方 \text{DFT}^2

\begin{aligned} (\text{DFT}^2)_{i,j}&=\sum_{k=0}^{2^m-1}\omega_{2^m}^{ik+kj}\\ &=\sum_{k=0}^{2^m-1}(\omega_{2^m}^{i+j})^k\\ &=2^m[2^m\mid(i+j)] \end{aligned}

这个矩阵的效果等价于全乘 2^m 再反转 1\sim2^m-1 项。因此 \text{DFT}^{-1}=\text{DFT}\times\text{DFT}^{-2},效果就是先 DFT 再全除 2^m,再反转 1\sim2^m-1 项。

如何在模意义下表示复数呢?实际上,当模数取 998244353 时,有原根 3,而 998244353=119\times 2^{23}+1,则可以取 \omega_{2^m}^i=3^{\frac{998244352}{2^m}\times i}

接下来是一些转置原理技巧。设 \text{DFT}=M\times pp 代表位反转变换,M 代表后续线性变换。由于 \text{DFT}=\text{DFT}^T那么实际上两个序列的卷积就是:

\begin{aligned} H&=\text{DFT}^{-2}\times\text{DFT}\times((\text{DFT}\times F)\circ(\text{DFT}\times G)) \\ &=\text{DFT}^{-2}\times\text{DFT}\times((\text{DFT}^T\times F)\circ(\text{DFT}^T\times G)) \\ &=\text{DFT}^{-2}\times M\times p\times((p^T\times M^T\times F)\circ(p^T\times M^T\times G)) \\ &=\text{DFT}^{-2}\times M\times((M^T\times F)\circ(M^T\times G)) \end{aligned}

由于 p^Tp 是互为逆的一对置换,在进行逐元素乘积前后可以直接抵消。因此只需得出 M 的转置算法即可,可以直接省掉位反转变换。转置后需要反转为自顶向下执行算法,最初一层内一对元素执行的操作是

\begin{bmatrix} f_{k}\\ f_{k+2^{m-1}} \end{bmatrix}= \begin{bmatrix} 1&\omega_{2^m}^k\\ 1&-\omega_{2^m}^{k} \end{bmatrix} \begin{bmatrix} f_{k}\\ f_{k+2^{m-1}} \end{bmatrix}

转置之后的矩阵就是

\begin{bmatrix} f_{k}\\ f_{k+2^{m-1}} \end{bmatrix}= \begin{bmatrix} 1&1\\ \omega_{2^m}^k&-\omega_{2^m}^{k} \end{bmatrix} \begin{bmatrix} f_{k}\\ f_{k+2^{m-1}} \end{bmatrix}

发现两者分别是

\begin{bmatrix} 1&1\\ 1&-1 \end{bmatrix} \begin{bmatrix} 1&0\\ 0&\omega_{2^m}^{k} \end{bmatrix}

\begin{bmatrix} 1&0\\ 0&\omega_{2^m}^{k} \end{bmatrix} \begin{bmatrix} 1&1\\ 1&-1 \end{bmatrix}

那么可以写出这样的代码:

#include <iostream>

const int P = 998244353, O=1<<21;
int N,M,A,i,j,V[O+1],W[O+1];
long long w,c,Z,R;
#define p a[j]
#define q a[j+i]
void G(int a[], int t) {
  for (unsigned I = t < 0 ? 0 : 20; I < 21; I -= t) {
    i = 1 << I;
    Z=3,w=1,R=P>>I+1;
    while(R)R>>=!(R & 1&&!(w=w*Z%P)),Z=Z*Z%P;
    for (j = 0, c = 1; j < O; c=++j&i?!!(j+=i):c*w%P) 
      (t < 0) && (q=1LL*q*c%P),
      p += q, q=p-q-q+P,
      q= 1LL*q*(t>0?c:1)%P,
      p%=P;
  }
}
main() {
  scanf("%d%d",&N,&M);
  for (i = 0; i <= N; ++i) scanf("%d",V+i);
  for (i = 0; i <= M; ++i) scanf("%d",W+i);
  G(V, 1), G(W, 1);
  for (i = 0; i < O; ++i) V[i] = c * V[i] * W[i] % P;
  G(V, -1);
  for (i = 0; i <= N + M; ++i) printf("%lld ", V[(O - i) & (O - 1)] * (c * P - P/O) % P);
}

然后就是大力 golf 了,展示最终成果。这里 15311432\equiv3^{119}

#import<cstdio>
const int O=1<<21,P=476*O+1;long w,c=1,N,M,A,*j,V[O],W[O],t,Z;void G(long*a){for(auto I=20u*!t;I<21;I+=t?:-1){for(w=15311432,Z=22-I;Z--;w=w*w%P)A=1<<I;for(j=a;j<a+O;c=++j-a&A?!!(j+=A):c*w%P)j[A]=((*j+=j[A]=j[A]%P*(t?c:1)%P)-j[A]*2+P)*(t?:c)%P,*j%=P;}}main(){for(scanf("%ld%ld",&N,&M),t=++N;t;--t)scanf("%ld",V+t);for(N+=M;~M;--M)scanf("%ld",W+M);for(G(V),G(W);t<O;++t)V[t]*=W[t];for(t=1,G(V);N;--N)printf("%ld ",V[O-N]*(P-P/O)%P);}