SPOJ Triple Sums(FFT+容斥原理)

Posted 2020-10-03 freeloop

 

# include <cstdio>
# include <cstring>
# include <cstdlib>
# include <iostream>
# include <vector>
# include <queue>
# include <stack>
# include <map>
# include <complex>
# include <set>
# include <cmath>
# include <algorithm>
using namespace std;
# define lowbit(x) ((x)&(-x))
const double pi=acos(-1.0);
# define eps 1e-8
# define MOD 1000000007
# define INF 1000000000
# define mem(a,b) memset(a,b,sizeof(a))
# define FOR(i,a,n) for(int i=a; i<=n; ++i)
# define FDR(i,a,n) for(int i=a; i>=n; --i)
# define bug puts("H");
# define lch p<<1,l,mid
# define rch p<<1|1,mid+1,r
# define mp make_pair
# define pb push_back
typedef pair<int,int> PII;
typedef vector<int> VI;
# pragma comment(linker, "/STACK:1024000000,1024000000")
typedef long long LL;
inline int Scan() {
    int x=0,f=1; char ch=getchar();
    while(ch<\'0\'||ch>\'9\'){if(ch==\'-\') f=-1; ch=getchar();}
    while(ch>=\'0\'&&ch<=\'9\'){x=x*10+ch-\'0\'; ch=getchar();}
    return x*f;
}
inline void Out(int a) {
    if(a<0) {putchar(\'-\'); a=-a;}
    if(a>=10) Out(a/10);
    putchar(a%10+\'0\');
}
const int N=80005;
//Code begin....

typedef complex<double> cmx;
int c[N<<2], val[N<<2], a[N<<2], b[N<<2];
cmx x[N<<2], y[N<<2];

void change(cmx x[], int len) {
    int i, j, k;
    for(i=1, j=len>>1; i<len-1; ++i) {
        if(i<j) swap(x[i],x[j]);
        k=len>>1;
        while(j>=k) j-=k, k>>=1;
        if(j<k) j+=k;
    }
}
void fft(cmx x[], int len, int on) {
    change(x,len);
    for(int i=2; i<=len; i<<=1) {
        cmx wn(cos(-on*2*pi/i),sin(-on*2*pi/i));
        for(int j=0; j<len; j+=i) {
            cmx w(1,0);
            FOR(k,j,j+i/2-1) {
                cmx u=x[k], v=x[k+i/2]*w;
                x[k]=u+v; x[k+i/2]=u-v; w*=wn;
            }
        }
    }
    if(on==-1) FOR(i,0,len-1) x[i]/=len;
}
int main()
{
    int n, maxx=0;
    cmx three(3.0,0);
    scanf("%d",&n);
    FOR(i,1,n) scanf("%d",val+i), val[i]+=20000, maxx=max(maxx,val[i]);
    maxx*=3;
    FOR(i,1,n) ++a[val[i]], ++b[val[i]*2], ++c[val[i]*3];
    int len=1;
    while(len<maxx) len<<=1; //len要为2的幂次
    FOR(i,0,len-1) x[i]=cmx(a[i],0), y[i]=cmx(b[i],0); 
    fft(x,len,1); fft(y,len,1); //将系数表达式转化为点值表达式
    FOR(i,0,len-1) x[i]=x[i]*x[i]*x[i]-x[i]*y[i]*three;
    fft(x,len,-1);
    FOR(i,0,len-1) a[i]=(int(x[i].real()+0.5)+2*c[i])/6;
    FOR(i,0,len-1) {
        if(!a[i]) continue;
        printf("%d : %d\\n", i-3*20000,a[i]);
    }
    return 0;
}

View Code