Triple Sums (FFT+容斥)

Triple Sums (FFT+容斥)

题目地址:洛谷 SP8372 TSUM - Triple Sums

首先构造函数\(f(x)=\sum\limits_{i=1}^{n}A_i=\sum\limits_{i}^\infty a_i*i​\),那么\(a_i\)表示的就是i是否出现过

那么题意就是求\(ans_k=\sum\limits_{i}^{k}a_i\sum\limits_{x}^{k-i}a_x a_{k-i-x}[i<x<k-i-x]\)

首先考虑可以重复选同一个数怎么做

那么先两遍FFT乘起来

然后去重,由于三个数有3!种排列方式,所以我们只需把答案除以6就好了.

那如果不能选相同的数怎么办?

考虑容斥去重

设\(b_i=a_{\frac{i}{2}}\)表示同一个数加两次

设\(c_i=a_{\frac{i}{3}}\)表示同一个数加三次

先把a,b,c都DFT一遍

则有\(d_i=a_i^3-3a_ib_i+2c_i\)表示三个不同的数加起来

然后把d给IDFT回去再除6,就是ans了

code:

#pragma GCC optimize(3)
#include<bits/stdc++.h>
#define N (1<<17)
#define offset 20000
#define pi 3.1415926535
#define double long double
#define int long long
using namespace std;
struct Complex{double x,y;};
Complex operator + (Complex a,Complex b){return (Complex){a.x+b.x,a.y+b.y};}
Complex operator - (Complex a,Complex b){return (Complex){a.x-b.x,a.y-b.y};}
Complex operator * (Complex a,Complex b){return (Complex){a.x*b.x-a.y*b.y,a.y*b.x+a.x*b.y};}
Complex a[N],b[N],c[N],d[N];
int rota[N],cnt;
void pre(int n){
    int high=0;
    cnt=1;
    while(cnt<=n)cnt<<=1,high++;
    for(int i=0;i<cnt;i++)
        rota[i]=(rota[i>>1]>>1)|((i&1)<<(high-1));
}
void fft(int lim,Complex *buf,int dft){
    for(int i=0;i<lim;i++)if(i<rota[i])swap(buf[i],buf[rota[i]]);
    for(int len=2;len<=lim;len<<=1){
        Complex wn=(Complex){cos(dft*2*pi/len),sin(dft*2*pi/len)};
        for(int s=0;s<lim;s+=len){
            Complex w=(Complex){1,0};
            for(int k=s;k<s+len/2;k++,w=w*wn){
                Complex x,y;
                x=buf[k],y=w*buf[k+len/2];
                buf[k]=x+y;
                buf[k+len/2]=x-y;
            }
        }
    }
    if(dft==-1)for(int i=0;i<lim;i++)buf[i].x/=lim;
}
int n,m,mx;
signed main(){
    scanf("%lld",&n);
    for(int i=0,k;i<n;i++){
        scanf("%lld",&k);
        k+=offset;
        a[k].x++;
        b[2*k].x++;
        c[3*k].x++;
    }
    pre((1<<17)-1);
    fft(cnt,a,1);
    fft(cnt,b,1);
    fft(cnt,c,1);
    for(int i=0;i<cnt;i++)
        d[i]=a[i]*(a[i]*a[i]-b[i]-b[i]-b[i])+(Complex){2,0}*c[i];
    fft(cnt,d,-1);
    for(int i=0;i<cnt;i++){
        long long ans=(long long)(d[i].x+0.5)/6;
        if(ans>0)printf("%lld : %lld\n",i-3*offset,ans);
    }
    return 0;
}
/*
5 
-1 
2 
3 
0 
5 
 
*/