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;
}