分塊FFT--bzoj3509: [CodeChef] COUNTARI
阿新 • • 發佈:2018-12-02
給定一個長度為N的陣列A[],求有多少對i, j, k(1<=i<j<k<=N)滿足A[k]-A[j]=A[j]-A[i]。
暴力的話就考慮列舉 ,然後算兩邊的可以用 來算,但這樣複雜度是 ,優化的話就可以用分塊,而且塊不能太小,然後塊外用 ,塊內就暴力。
注意卡常
被卡了無數次以後發現不可以預處理要邊做邊處理
#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#define maxn 30005
#define maxm 100005
using namespace std;
const double Pi=acos(-1.0);
inline int rd(){
int x=0,f=1;char c=' ';
while(c<'0' || c>'9') f=c=='-'?-1:1,c=getchar();
while(c<='9' && c>='0') x=x*10+c-'0',c=getchar();
return x*f;
}
inline int max(int x,int y){return x>y?x:y;}
inline int min(int x,int y){return x<y?x:y;}
struct complex{
double x,y;
complex(double xx=0,double yy=0) {x=xx,y=yy;}
}a[maxn<<2],b[maxn<<2];
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.x*b.y+a.y*b.x);}
int n,cntl[maxn],cntr[maxn],limit=1,l,rev[maxn<<2];
int block,L[300],R[300],tot,mx,val[maxm];
long long res;
inline void FFT(complex *F,int type){
for(int i=0;i<limit;i++) if(i<rev[i]) swap(F[i],F[rev[i]]);
for(int mid=1;mid<limit;mid<<=1){
complex Wn(cos(Pi/mid),1.0*type*sin(Pi/mid));
for(int r=mid<<1,j=0;j<limit;j+=r){
complex w(1,0);
for(int k=0;k<mid;k++,w=w*Wn){
complex x=F[j+k],y=w*F[j+mid+k];
F[j+k]=x+y,F[j+mid+k]=x-y;
}
}
}
if(type==-1){
for(int i=0;i<limit;i++) F[i].x=F[i].x/limit;
}
}
inline void prework(){
for(int i=1;i<=tot;i++) L[i]=R[i-1]+1,R[i]=i*block;
if(R[tot]<n) ++tot,L[tot]=R[tot-1]+1,R[tot]=n;
while(limit<=2*mx) limit<<=1,++l;
for(int i=0;i<limit;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
}
inline void solve(){
for(int i=1;i<=tot;i++){
for(int j=0;j<limit;j++) a[j].x=a[j].y=b[j].x=b[j].y=0;
for(int j=L[i];j<=R[i];j++) cntr[val[j]]--;
for(int j=0;j<=mx;j++) a[j].x=cntl[j],b[j].x=cntr[j];
FFT(a,1); FFT(b,1);
for(int j=0;j<limit;j++) a[j]=a[j]*b[j];
FFT(a,-1);
for(int j=L[i];j<=R[i];j++) {
res+=(int)(a[val[j]<<1].x+0.5);
for(int k=L[i];k<=R[i];k++){
if(k<j){
int x=2*val[j]-val[k];
if(x>=0 && x<=mx) res+=cntr[x];
}
else if(k>j){
int x=2*val[j]-val[k];
if(x>=0 && x<=mx) res+=cntl[x];
}
}
cntl[val[j]]++;
}
}
}
int main(){
n=rd();
for(int i=1;i<=n;i++) val[i]=rd(),mx=max(mx,val[i]),cntr[val[i]]++;
block=min((int)6*sqrt(n),n); tot=n/block;
prework(); solve();
printf("%lld\n",res);
return 0;
}