bzoj 3771: Triple （容斥原理+生成函数+FFT）

题目描述

传送门

题目大意：有价值不同的n个物件，求从中选取1-3件物品，能组成多少种不同的价值，已经每种价值的方案数。

题解

容斥原理+生成函数+FFT
因为情况较少，所以可以手动容斥一下。
设每个物品选取1件得到的生成函数序列为a
每个物品选取2件得到的生成函数序列为b
每个物品选取3件得到的生成函数序列为c
那么选走1个物品得到的生成函数的序列就是a
选走2个物品得到的生成函数的序列就是$(a*a-b)/2$
选走3个物品得到的生成函数的序列就是$(a*a*a-a*b*3+2*c)/6$

代码

#include<iostream>
#include<cstring>
#include<algorithm>
#include<cstdio>
#include<cmath>
#define N 200003
#define pi acos(-1)
using namespace std;
int n,m,n1,val[N],L,R[N],ans[N],ans1[N],ans0[N];
struct data{
    double x,y; 
    data(double X=0,double Y=0) {
        x=X,y=Y;
    }
}f[N],f1[N],g[N],h[N];
data operator +(data a,data b){
    return data(a.x+b.x,a.y+b.y);
}
data operator -(data a,data b){
    return data(a.x-b.x,a.y-b.y);
}
data operator *(data a,data b){
    return data(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);
}
void FFT(data a[N],int n,int opt)
{
    for (int i=0;i<n;i++)
     if (i>R[i]) swap(a[i],a[R[i]]);
    for (int i=1;i<n;i<<=1) {
        data wn=data(cos(pi/i),opt*sin(pi/i));
        for (int p=i<<1,j=0;j<n;j+=p) {
            data w=data(1,0);
            for (int k=0;k<i;k++,w=w*wn){
                data x=a[j+k]; data y=a[j+k+i]*w;
                a[j+k]=x+y; a[j+k+i]=x-y;
            }
        }
    }
    if (opt==-1)
     for (int i=0;i<n;i++) a[i].x/=n;
}
int main()
{
    freopen("a.in","r",stdin);
    freopen("my.out","w",stdout);
    scanf("%d",&n); int mx=0;
    for (int i=1;i<=n;i++) scanf("%d",&val[i]),mx=max(mx,val[i]);
    for (int i=1;i<=n;i++) 
     f[val[i]].x=1,f1[val[i]*2].x=1,
     ans0[val[i]]=1,ans[val[i]*2]-=1,ans1[val[i]*3]+=2;
    m=3*mx;
    for (n1=1;n1<=m;n1<<=1) L++;
    for (int i=0;i<=n1;i++) R[i]=(R[i>>1]>>1)|((i&1)<<(L-1));
    FFT(f,n1,1); FFT(f1,n1,1);
    for (int i=0;i<=n1;i++) g[i]=f[i]*f[i],h[i]=g[i];
    FFT(g,n1,-1);
    for (int i=0;i<=n1;i++) ans[i]+=(int)(g[i].x+0.1);
    for (int i=0;i<=n1;i++) h[i]=h[i]*f[i];
    FFT(h,n1,-1);
    for (int i=0;i<=n1;i++) ans1[i]+=(int)(h[i].x+0.1);
    for (int i=0;i<=n1;i++) h[i]=f[i]*f1[i];
    FFT(h,n1,-1);
    for (int i=0;i<=n1;i++) ans1[i]-=(int)(h[i].x+0.1)*3;
    for (int i=1;i<=mx*3;i++) ans0[i]+=ans[i]/2+ans1[i]/6;
    for (int i=1;i<=mx*3;i++)
     if (ans0[i]) printf("%d %d\n",i,ans0[i]);
    return 0;
}