BZOJ3771

Portal

Description

给出(n)个互不相同的数({a_n}(a_ileq4 imes10^4))，从中选出一至三个数，其和为(sum)。求对于所有可能的(sum)各有多少种方案。

Solution

生成函数+容斥原理+FFT。
定义(A(x))为仅选取一个数的生成函数，即(A(a_i)=1)；(B(x))为仅选取两个相同数的生成函数，即(B(2a_i)=1)；(C(x))为仅选取三个相同数的生成函数，即(C(3a_i)=1)。

仅选取一个数的生成函数为(A)。
仅选取两个数的生成函数为((A*A-B)/2)。
仅选取三个数的生成函数为((A*A*A-3A*B+2C)/6)。

依容斥原理简单解释一下。选取两个数：等价于随便选两个数(x_1,x_2)（(A*A)），减去(x_1=x_2)（(B)），因为选择有顺序所以除以(2!)。选取三个数：等价于随便选三个数(x_1,x_2,x_3)（(A*A*A)），减去(x_1=x_2)和(x_2=x_3)和(x_3=x_1)（(3A*B)），加上两倍（因为被减三次）(x_1=x_2=x_3)，再除以(3!)。
求卷积的话用FFT就好。

设(n_0=3max{a})，时间复杂度为(O(n_0logn_0))。

Code

//Triple
#include <complex>
#include <cstdio>
typedef long long lint;
typedef std::complex<double> cpx;
const int N=15e4;
const double PI=acos(-1);
int n;
int n0,t,pos[N];
cpx a[N],b[N],c[N],ans[N];
int max(int x,int y) {return x>y?x:y;}
void FFT(cpx x[],int f)
{
    for(int i=0;i<n0;i++) if(i<pos[i]) swap(x[i],x[pos[i]]);
    for(int i=1;i<n0;i<<=1)
    {
        cpx Wn=cpx(cos(PI/i),f*sin(PI/i));
        for(int j=0;j<n0;j+=i<<1)
        {
            cpx w=1;
            for(int k=0;k<i;k++,w*=Wn)
            {
                cpx p=x[j+k],q=w*x[j+k+i];
                x[j+k]=p+q,x[j+k+i]=p-q;
            }
        }
    }
    if(f==-1) for(int i=0;i<n0;i++) x[i]/=n0;
}
int main()
{
    scanf("%d",&n); int m=0;
    for(int i=1;i<=n;i++)
    {
        int x; scanf("%d",&x);
        m=max(m,x);
        a[x]=b[x+x]=c[x+x+x]=1;
    }
    n0=1,t=0; while(n0<m+m+m) n0<<=1,t++;
    for(int i=0;i<n0;i++) pos[i]=pos[i>>1]>>1|(i&1)<<t-1;
    FFT(a,1),FFT(b,1),FFT(c,1);
    for(int i=0;i<n0;i++)
    {
        ans[i]=a[i];
        ans[i]+=(a[i]*a[i]-b[i])/2.0;
        ans[i]+=(a[i]*a[i]*a[i]-3.0*a[i]*b[i]+2.0*c[i])/6.0;
    }
    FFT(ans,-1);
    for(int i=0;i<n0;i++)
    {
        lint x=(lint)(ans[i].real()+0.5);
        if(x) printf("%d %lld
",i,x);
    }
    return 0;
}