NTT 模板

Posted 2020-12-01 bakacirno

其实和 (FFT) 就是一个模子里刻出来的一样,
(FFT) 的优化是基于它的单位根
而 (NTT) 的模数通常有一个原根, 和 (FFT) 的单位根有类似的性质
还是存个模板
模板题

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#define MAXN 3000010
using namespace std;
typedef long long LL;
const int mod=998244353;
const int g=3;
int n,m,rev[MAXN],bit=0,len=1;
LL a[MAXN],b[MAXN];
int read(){
    int x=0,f=1;char c=getchar();
    while(c<'0'||c>'9') {if(c=='-')f=-1;c=getchar();}
    while(c>='0'&&c<='9') {x=x*10+c-'0';c=getchar();}
    return x*f;
}

LL qpow(LL x,LL k){
    LL res=1;
    while(k){
        if(k&1) res=res*x%mod;
        x=x*x%mod;
        k>>=1;
    }
    return res;
}

void NTT(LL *a,int opt){
    for(int i=0;i<len;i++)
        if(i<rev[i]) swap(a[i],a[rev[i]]);
    for(int mid=1;mid<len;mid<<=1){
        LL tmp=qpow(g,(mod-1)/(mid*2));
        if(opt==-1) tmp=qpow(tmp,mod-2);
        for(int i=0;i<len;i+=mid*2){
            LL w=1;
            for(int j=0;j<mid;j++,w=w*tmp%mod){
                LL x=a[i+j],y=w*a[i+j+mid]%mod;
                a[i+j]=(x+y)%mod,a[i+j+mid]=(x-y+mod)%mod;
            }
        }
    }
}

int main(){
    n=read();m=read();
    for(int i=0;i<=n;i++) a[i]=read();
    for(int i=0;i<=m;i++) b[i]=read();
    while(len<=n+m) len<<=1,bit++;
    for(int i=0;i<len;i++) rev[i]=(rev[i>>1]>>1) | ((i&1)<<(bit-1));
    NTT(a,1);NTT(b,1);
    for(int i=0;i<len;i++) a[i]=a[i]*b[i]%mod;
    NTT(a,-1);
    LL inv=qpow(len,mod-2);
    for(int i=0;i<=n+m;i++) printf("%lld ",a[i]*inv%mod);
    return 0;
}