题目地址:洛谷 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}\)
首先考虑可以重复选同一个数怎么做
那么先两遍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
*/