首页 > 代码库 > NTT+多项式求逆+多项式开方(BZOJ3625)

NTT+多项式求逆+多项式开方(BZOJ3625)

定义多项式h(x)的每一项系数hi,为i在c[1]~c[n]中的出现次数。

定义多项式f(x)的每一项系数fi,为权值为i的方案数。

通过简单的分析我们可以发现:f(x)=2/(sqrt(1-4h(x))+1)

于是我们需要多项式开方和多项式求逆。

多项式求逆:

求B(x),使得A(x)*B(x)=1 (mod x^m)

考虑倍增。

假设我们已知A(x)*B(x)=1 (mod x^m),要求C(x),使得A(x)*C(x)=1 (mod x^(2m))

简单分析可得C(x)=B(x)*(2-A(x)*B(x))

多项式开方:

求B(x),使得B(x)*B(x)=A(x) (mod x^m)

继续考虑倍增。

假设我们已知B(x)*B(x)=A(x) (mod x^m),要求C(x),使得C(x)*C(x)=A(x) (mod x^(2m))

简单分析可得C(x)=(B(x)^2+A(x))/(2B(x)),需要求B(x)的逆。

观察到以上两个式子都用到了多项式乘法,又因为在模意义下,我们需要NTT。

NTT和FFT的区别在于单位根不一样,NTT需要找一个p的原根,而且通常要求p是形如k*2^q+1的一个质数。

找原根暴力找就行了。

(这题卡常数,所以我都用的int)

#include <cstdio>
#include <cstring>
#include <algorithm>

typedef long long ll;
const int N=300000,p=998244353;
int n,m,x,ni,r[N],h[N],_h[N],__h[N],t[N];
ll pw(ll a,int b) {
    ll r=1;
    for(;b;b>>=1,a=a*a%p) if(b&1) r=r*a%p;
    return r;
}

void ntt(int *a,int n,int f) {
    for(int i=0;i<n;i++) if(r[i]>i) std::swap(a[i],a[r[i]]);
    for(int i=2;i<=n;i<<=1)
    for(int j=0,wn=pw(3,((p-1)/i*f+p-1)%(p-1)),m=i>>1;j<n;j+=i)
    for(int k=0,w=1;k<m;k++,w=(ll)w*wn%p) {
        int x=a[j+k],y=(ll)a[j+k+m]*w%p;
        a[j+k]=(x+y)%p,a[j+k+m]=(x-y+p)%p;
    }
    if(f==-1) {
        ll ni=pw(n,p-2);
        for(int i=0;i<n;i++) a[i]=(ll)a[i]*ni%p;
    }
}
void gt2(int n) {
    if(n==1) {__h[0]=pw(_h[0],p-2); return;}
    gt2(n>>1),memcpy(t,_h,sizeof(int)*n),memset(t+n,0,sizeof(int)*n);
    int m=1,l=-1;
    while(m<n<<1) m<<=1,l++;
    for(int i=0;i<m;i++) r[i]=(r[i>>1]>>1)|((i&1)<<l);
    ntt(t,m,1),ntt(__h,m,1);
    for(int i=0;i<m;i++) __h[i]=(ll)__h[i]*(2-(ll)t[i]*__h[i]%p+p)%p;
    ntt(__h,m,-1),memset(__h+n,0,sizeof(int)*n);
}
void gt(int n) {
    if(n==1) {_h[0]=1; return;}
    gt(n>>1),memset(__h,0,sizeof(int)*n),gt2(n);
    memcpy(t,h,sizeof(int)*n),memset(t+n,0,sizeof(int)*n);
    int l=-1,m=1;
    while(m<n<<1) m<<=1,l++;
    for(int i=0;i<m;i++) r[i]=(r[i>>1]>>1)|((i&1)<<l);
    ntt(t,m,1),ntt(_h,m,1),ntt(__h,m,1);
    for(int i=0;i<m;i++) _h[i]=((ll)_h[i]*_h[i]+t[i])%p*__h[i]%p*ni%p;
    ntt(_h,m,-1),memset(_h+n,0,sizeof(int)*n);
}

int main() {
    scanf("%d%d",&m,&n),ni=pw(2,p-2);
    for(int i=1;i<=m;i++) scanf("%d",&x),h[x]++;
    for(int i=1;i<=n;i++) h[i]=(-h[i]*4+p)%p;
    for(m=n,n=1;n<=m;n<<=1);
    h[0]++,gt(n),_h[0]=(_h[0]+1)%p,memset(__h,0,sizeof __h),gt2(n);
    for(int i=1;i<=m;i++) printf("%d\n",__h[i]*2%p);
    return 0;
}

NTT+多项式求逆+多项式开方(BZOJ3625)