BZOJ 3771: Triple（生成函数+FFT+容斥）

题目描述

传送门

题目大意：给一堆不同的数，问你从中取1-3个组成不同和的方案数，不考虑顺序。每个数的大小<=40000。

---

思路

简单的容斥。搞出取一个的生成函数$A$、取两个相同的生成函数$B$、取三个相同的生成函数$C$，然后答案就是$A+A^2-B+A^3-3AB+2C$。

乘起来的过程用FFT加速就行了。这是考虑了顺序的，不考虑顺序的话在加上的时候除以对应的阶乘就行了。

ps：一开始次数界搞成了非二的幂，死活找不到错的感觉。

---

代码

#include <bits/stdc++.h>
#define MAXN 120010

using namespace std;

typedef long long LL;
const double PI = acos(-1.0);
int n;

struct Complex{
    double real, image;
    Complex() {}
    Complex(double _real, double _image){
        real = _real;
        image = _image;
    }
    friend Complex operator + (Complex A, Complex B){
        return Complex(A.real + B.real, A.image + B.image);
    }
    friend Complex operator - (Complex A, Complex B){
        return Complex(A.real - B.real, A.image - B.image);
    }
    friend Complex operator * (Complex A, Complex B){
        return Complex(A.real * B.real - A.image * B.image, A.real * B.image + A.image * B.real);
    }
}a[MAXN<<2], b[MAXN<<2], c[MAXN<<2], d[MAXN<<2];

LL ans[MAXN<<2];

int Get(int x){
    int t = 1;
    while(t < x)  t <<= 1;
    t <<= 1;
    return t;
}

void Reverse(Complex *A){
    for(int i = 0; i < n-1; i++){
        int j = 0;
        for(int k = 1, tmp = i; k < n; k <<= 1, tmp >>= 1)
            j = ((j << 1) | (tmp & 1));
        if(j > i)  swap(A[i], A[j]); 
    }
}

void FFT(Complex *A, int DFT){
    Reverse(A);
    for(int s = 1; (1 << s) <= n; s++){
        int m = 1 << s;
        Complex wn = Complex(cos(2*PI*DFT/m), sin(2*PI*DFT/m));
        for(int k = 0; k < n; k += m){
            Complex w = Complex(1, 0);
            for(int j = 0; j < (m>>1); j++){
                Complex u = A[k+j], t = w * A[k+j+(m>>1)];          
                A[k+j] = u + t;
                A[k+j+(m>>1)] = u - t;
                w = w * wn;
            }
        }
    }
    if(DFT == -1)  for(int i = 0; i < n; i++)  A[i].real /= n;
}

int main(){

    scanf("%d", &n);

    int Max = 0, x;
    for(int i = 1; i <= n; i++){
        scanf("%d", &x);
        a[x].real += 1.0;
        b[x+x].real += 1.0;
        c[x+x+x].real += 1.0;
        Max = max(Max, x);
    }

    Max *= 3;
    n = Get(Max+2);

    for(int i = 0; i < n; i++)  ans[i] += (LL)(a[i].real + .5);

    FFT(a, 1);  for(int i = 0; i < n; i++)  d[i] = a[i] * a[i];  FFT(d, -1);

    for(int i = 0; i < n; i++)  ans[i] += ((LL) (d[i].real + .5) - (LL)(b[i].real + .5)) / 2;

    for(int i = 0; i < n; i++)  d[i] = a[i] * a[i] * a[i];  FFT(d, -1);
    FFT(b, 1);  for(int i = 0; i < n; i++)  b[i] = a[i] * b[i];  FFT(b, -1);

    for(int i = 0; i < n; i++)  
        ans[i] += ((LL)(d[i].real + .5) - 3 * (LL)(b[i].real + .5) + 2 * (LL)(c[i].real + .5)) / 6;

    for(int i = 0; i <= Max; i++)
        if(ans[i])  printf("%d %lld\n", i, ans[i]);

    return 0;
}

---

此去泉台招旧部，旌旗十万斩阎罗。