BZOJ 3771 Triple (FFT+生成函数+容斥)

题面传送门

题目大意：给你互不相同的$n$个数，在其中任选$1$~$3$个数，不能重复选数，设它们的和为$S$。对于所有可能的$S$，求选出的数和为$S$方案总数，选数没有顺序。

先对所有的数弄一个生成函数$A$，有数的位置权值为$1$

如果我们要选$x$个数，求方案数。只需要对$A$求$x$次卷积。

然而题目要求选数是无序的，且不能重复选数

选$1$个数的方案就是$A$

选$2$个数时，需要去掉重复选同一个数的情况，即$(A*A-B)/2$，$B$是相同的数选$2$次时的生成函数

选$3$个数的情况更复杂

1.首先要去掉重复选$3$次同一个数的情况，设相同的数选$3$次时的生成函数为$C$

2.还要去掉 选$2$次相同的数，再选另一个数的情况，这种情况的生成函数是$3*B*A$

然而情况二还可能把情况一去掉，需要再加回来

最终答案就是$A+(A*A-B)/2+(A*A*A-3*B*A+2*C)/6$

#include <cmath>
#include <cstdio>
#include <cstring>
#include <algorithm>
#define N1 (1<<18)+10
#define ll long long 
#define dd double
using namespace std;
 
 
int gint()
{
    int ret=0,fh=1; char c=getchar();
    while(c<'0'||c>'9'){if(c=='-')fh=-1;c=getchar();}
    while(c>='0'&&c<='9'){ret=ret*10+c-'0';c=getchar();}
    return ret*fh;
}
 
int r[N1];
const dd pi=acos(-1);
struct cp{
dd x,y;
friend cp operator + (const cp &s1,const cp &s2){ return (cp){s1.x+s2.x,s1.y+s2.y}; }
friend cp operator - (const cp &s1,const cp &s2){ return (cp){s1.x-s2.x,s1.y-s2.y}; }
friend cp operator * (const cp &s1,const cp &s2){ return (cp){s1.x*s2.x-s1.y*s2.y,s1.y*s2.x+s1.x*s2.y}; }
}a[N1],b[N1],c[N1],f[N1];
 
void init()
{
    memset(a,0,sizeof(a));
    memset(b,0,sizeof(b));
    memset(c,0,sizeof(c));
}
 
void FFT(cp *s,int len,int type)
{
    int i,j,k; cp wn,w,t;
    for(i=0;i<len;i++) if(i<r[i]) swap(s[i],s[r[i]]);
    for(k=2;k<=len;k<<=1)
    {
        wn=(cp){cos(2.0*pi*type/k),sin(2.0*pi*type/k)};
        for(i=0;i<len;i+=k)
        {
            w=(cp){1,0};
            for(j=0;j<(k>>1);j++,w=w*wn)
            {
                t=w*s[i+j+(k>>1)];
                s[i+j+(k>>1)]=s[i+j]-t;
                s[i+j]=s[i+j]+t;
            }
        }
    }
}
int n,ma,len,L;
int v[N1],s1[N1],s2[N1],s3[N1],ans[N1];
 
int main()
{
    scanf("%d",&n); int i;
    for(i=1;i<=n;i++) v[i]=gint(),ma=max(ma,v[i]),s1[v[i]]++; // 1 -> A
     
    for(len=1,L=0;len<3*ma;len<<=1,L++);
    for(i=1;i<=n;i++) a[v[i]].x=1,b[v[i]<<1].x=1;//c[v[i]*3].x=1;
    for(i=0;i<len;i++) r[i]=(r[i>>1]>>1)|((i&1)<<(L-1));
    FFT(a,len,1); FFT(b,len,1);// FFT(c,len,1);
     
    for(i=0;i<len;i++) f[i]=a[i]*a[i]; // 2 -> (A*A-B)/2
    FFT(f,len,-1);
    for(i=1;i<=n;i++) s2[v[i]<<1]--;
    for(i=0;i<len;i++) s2[i]=(int)(f[i].x/len+0.1);
    for(i=0;i<len;i++) s2[i]/=2;
     
    for(i=0;i<len;i++) f[i]=a[i]*a[i]*a[i]; // 3 -> (A*A*A-3*B*A+2*C)/6
    FFT(f,len,-1);
    for(i=0;i<len;i++) s3[i]+=(int)(f[i].x/len+0.1);
    for(i=0;i<len;i++) f[i]=a[i]*b[i];
    FFT(f,len,-1);
    for(i=0;i<len;i++) s3[i]-=( (int)(f[i].x/len+0.1) )*3;
    for(i=1;i<=n;i++) s3[v[i]*3]+=2;
    for(i=0;i<len;i++) s3[i]/=6;
     
    for(i=1;i<=len;i++) if(s1[i]+s2[i]+s3[i]>0) 
        printf("%d %d
",i,s1[i]+s2[i]+s3[i]);
     
    return 0;
}