BZOJ3771 - Triple

Portal

Description

给出\(n\)个互不相同的数\(\{a_n\}(a_i\leq4\times10^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\n",i,x);
    }
    return 0;
}