SPOJ - TSUM Triple Sums (FFT+容斥)

题意:给N个数,不重复地选3个数,求能够组成的数有多少种选法.
分析:若只选两个数就比较好求,FFT后减去两个相同的数构成的情况,再将每种情况除2(2个数排列有两种不同可能)即可.
选3个数也是类似地用容斥的方法计算,首先无限制地情况下,多项式中多计算了两个数相同和3个数相同的情况.有式:\(\sum(xyz) = \frac{A(x)^3 -3{B(x^2)*A(x)} + 2C(x^3)}{6}\).其中两个数相同的情况,排列有3种可能,三个数相同的情况,排列只一种可能.即多减去了2次三个数相同的情况,根据容斥原理,需要加回.

#include <bits/stdc++.h>
using namespace std;
using namespace std;
typedef long long LL;
const int MAXN = 4e5 + 10;
const double PI = acos(-1.0);
struct Complex{
    double x, y;
    inline Complex operator+(const Complex b) const {
        return (Complex){x +b.x,y + b.y};
    }
    inline Complex operator-(const Complex b) const {
        return (Complex){x -b.x,y - b.y};
    }
    inline Complex operator*(const Complex b) const {
        return (Complex){x *b.x -y * b.y,x * b.y + y * b.x};
    }
} va[MAXN * 2 + MAXN / 2], vb[MAXN * 2 + MAXN / 2];
int lenth = 1, rev[MAXN * 2 + MAXN / 2];
int N, M;   // f 和 g 的数量
    //f g和 的系数
    // 卷积结果
    // 大数乘积
int f[MAXN],g[MAXN];
vector<LL> conv;
vector<LL> multi;
//f g
void init()
{
    int tim = 0;
    lenth = 1;
    conv.clear(), multi.clear();
    memset(va, 0, sizeof va);
    memset(vb, 0, sizeof vb);
    while (lenth <= N + M - 2)
        lenth <<= 1, tim++;
    for (int i = 0; i < lenth; i++)
        rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (tim - 1));
}
void FFT(Complex *A, const int fla)
{
    for (int i = 0; i < lenth; i++){
        if (i < rev[i]){
            swap(A[i], A[rev[i]]);
        }
    }
    for (int i = 1; i < lenth; i <<= 1){
        const Complex w = (Complex){cos(PI / i), fla * sin(PI / i)};
        for (int j = 0; j < lenth; j += (i << 1)){
            Complex K = (Complex){1, 0};
            for (int k = 0; k < i; k++, K = K * w){
                const Complex x = A[j + k], y = K * A[j + k + i];
                A[j + k] = x + y;
                A[j + k + i] = x - y;
            }
        }
    }
}

int a[40005];
int cnt[MAXN],cnt2[MAXN],cnt3[MAXN];

void getConv(){             //求多项式
    init();
    for (int i = 0; i < N; i++)
        va[i].x = f[i];
    for (int i = 0; i < M; i++)
        vb[i].x = g[i];
    FFT(va, 1), FFT(vb, 1);
    for (int i = 0; i < lenth; i++)
        va[i] = va[i]*(va[i]*va[i] - (Complex){3.0,0.0}*vb[i]);
    FFT(va, -1);
    for (int i = 0; i <= N + M - 2; i++)
        conv.push_back(((LL)(va[i].x / lenth + 0.5)+2*cnt3[i])/6);
}

const int limit = 120005;

int main()
{
    #ifndef ONLINE_JUDGE
        freopen("in.txt","r",stdin);
        freopen("out.txt","w",stdout);
    #endif
    int n;
    scanf("%d",&n);
    int mx = -1;
    for(int i=1;i<=n;++i){
        scanf("%d",&a[i]);
        a[i]+= 20000;
        mx = max(mx,a[i]);
        cnt[a[i]]++;
        cnt2[a[i]*2]++;
        cnt3[a[i]*3]++;
    }
    N = mx+1;
    for(int i=0;i<limit;++i){
        f[i] = cnt[i];
    }
    M = mx*2+1;
    for(int i=0;i<limit;++i){
        g[i] = cnt2[i];
    }
    getConv();
    int sz = conv.size();
    for(int i=0;i<sz;++i){
        if(!conv[i]) continue;
        printf("%d : %lld\n",i-60000,conv[i]);
    }
    return 0;
}