BZOJ 4332: JSOI2012 分零食 FFT+分治

好题好题~

#include <bits/stdc++.h>

#define N 50020

#define ll long long

#define setIO(s) freopen(s".in","r",stdin)

using namespace std;

const double pi=acos(-1.0);

struct cpx

{

    double x,y;

    cpx(double a=0,double b=0){x=a,y=b; }

    cpx operator+(const cpx b) { return cpx(x+b.x,y+b.y);  }

    cpx operator-(const cpx b) { return cpx(x-b.x,y-b.y);  }

    cpx operator*(const cpx b) { return cpx(x*b.x-y*b.y,x*b.y+y*b.x); }

}a[N],b[N];

int pos[N],f[N],g[N],h[N],p[N],lim=1,m,mod;

void FFT(cpx *a,int len,int flag)

{

    int i,j,k,mid;

    for(i=k=0;i<len;++i)

    {

        if(i>k)   swap(a[i],a[k]);

        for(j=len>>1;(k^=j)<j;j>>=1);

    }

    for(mid=1;mid<len;mid<<=1)

    {

        cpx wn(cos(pi/mid), flag*sin(pi/mid)),x,y;

        for(i=0;i<len;i+=mid<<1)

        {

            cpx w(1,0);

            for(j=0;j<mid;++j)

            {

                x=a[i+j], y=w*a[i+j+mid];

                a[i+j]=x+y, a[i+j+mid]=x-y;

                w=w*wn;

            }

        }

    }

    if(flag==-1)    for(i=0;i<len;++i)    a[i].x/=(double)len;

}

inline void Mul(int x[],int y[],int z[])

{

    for(int i=0;i<lim;++i)      a[i]=cpx(x[i],0), b[i]=cpx(y[i],0);

    FFT(a,lim,1),FFT(b,lim,1);

    for(int i=0;i<lim;++i)    a[i]=a[i]*b[i];

    FFT(a,lim,-1);

    for(int i=0;i<=m;++i)     z[i]=(int)(floor(a[i].x+0.5))%mod;

}

void solve(int k)

{

    if(k==1)  return;

    solve(k/2);

    Mul(f,g,p),Mul(g,g,g);

    for(int i=0;i<lim;++i)   f[i]=(f[i]+p[i])%mod;

    if(k&1)

    {

        Mul(g,h,g);

        for(int i=0;i<lim;++i)   f[i]=(f[i]+g[i])%mod;

    }

    return;

}

int main()

{

    // setIO("input");

    int i,j,k,A,O,S,U;

    scanf("%d%d%d%d%d%d",&m,&mod,&A,&O,&S,&U);

    while(lim<=m<<1)    lim=lim<<1;

    for(i=1;i<=m;++i)   f[i]=g[i]=h[i]=(1ll*O*i%mod*i%mod+1ll*S*i%mod+U)%mod;

    solve(A),    printf("%d\n",f[m]);

    return 0;

}