Bzoj3509 [CodeChef] COUNTARI

Time Limit: 40 Sec Memory Limit: 128 MB
Submit: 912 Solved: 264

Description

给定一个长度为N的数组A[]，求有多少对i, j, k（1<=i<j<k<=N）满足A[k]-A[j]=A[j]-A[i]。

Input

第一行一个整数N（N<=10^5）。
接下来一行N个数A[i]（A[i]<=30000）。

Output

一行一个整数。

Sample Input

10
3 5 3 6 3 4 10 4 5 2

Sample Output

HINT

Source

分块+生成函数+FFT

看到A[k]-A[j]=A[j]-A[i]，下意识就移项得到了A[k]+A[i]=A[j]+A[j]

猛一看这不是普通的生成函数式子吗，枚举每个j位置，前后FFT卷积就解决了。

蛤蛤蛤，这题被我秒了——

这时博主眉头一皱，察觉事情有些不对，一算时间复杂度，是O(n*VlogV) Vmax=60000

（啪）我还是太年轻了

上面那个做法的缺点在于对于每个位置j算了它前后所有位置的卷积，未免太浪费了，如果可以算一次卷积就得到很多个j的答案，就很棒了。

于是考虑分块，块外FFT，块内暴力统计。O(n/m * VlogV)大概可以过。

写完不放心，网上找了篇题解的程序开始对拍。

发现大数据下拍不对，开始一点点查错，一个多小时过去了，无果。

怒把网上找的那程序交了一发，WA。

然后把自己的代码交上去，AC

F♂A♂Q

  1 /*by SilverN*/
  2 #include<algorithm>
  3 #include<iostream>
  4 #include<cstring>
  5 #include<cstdio>
  6 #include<cmath>
  7 #include<vector>
  8 #define LL long long
  9 using namespace std;
 10 const double pi=acos(-1.0);
 11 const int mxn=120010;
 12 int read(){
 13     int x=0,f=1;char ch=getchar();
 14     while(ch<'0' || ch>'9'){if(ch=='-')f=-1;ch=getchar();}
 15     while(ch>='0' && ch<='9'){x=x*10+ch-'0';ch=getchar();}
 16     return x*f;
 17 }
 18 //
 19 struct com{
 20     double x,y;
 21     com operator + (com b){return (com){x+b.x,y+b.y};}
 22     com operator - (com b){return (com){x-b.x,y-b.y};}
 23     com operator * (com b){return (com){x*b.x-y*b.y,x*b.y+y*b.x};}
 24 }a[mxn],b[mxn],c[mxn];
 25 int rev[mxn];
 26 int N,l;
 27 void R_init(int n){
 28     for(int i=0;i<n;i++)rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
 29     return;
 30 }
 31 void FFT(com *a,int flag){
 32     for(int i=0;i<N;i++)if(rev[i]>i)swap(a[rev[i]],a[i]);
 33     for(int i=1;i<N;i<<=1){
 34         com wn=(com){cos(pi/i),flag*sin(pi/i)};
 35         for(int j=0;j<N;j+=(i<<1)){
 36             com w=(com){1,0};
 37             for(int k=0;k<i;k++,w=w*wn){
 38                 com x=a[j+k],y=w*a[i+j+k];
 39                 a[j+k]=x+y;
 40                 a[i+j+k]=x-y;
 41             }
 42         }
 43     }
 44     if(flag==-1){for(int i=0;i<N;i++)a[i].x/=N;}
 45     return;
 46 }
 47 //FFT
 48 struct node{
 49     int L,R;
 50     int id;
 51 }B[mxn];
 52 int block,cnt;
 53 //分块 
 54 int w[mxn];
 55 LL ans=0;
 56 int Lc[mxn],Rc[mxn];
 57 int n;
 58 void solve(int id){
 59 //    printf("solve:%d
",id);
 60 //    printf("%lld
",ans);
 61     int i,j,x;
 62     for(i=B[id].L;i<=B[id].R;i++){//统计块内 
 63         --Rc[w[i]];
 64         for(j=B[id].L;j<i;j++){
 65             x=w[i]*2-w[j];
 66             if(x>=0)ans+=Rc[x];
 67         }
 68         for(j=i+1;j<=B[id].R;j++){
 69             x=w[i]*2-w[j];
 70             if(x>=0)ans+=Lc[x];
 71         }
 72     }
 73 //    return;
 74 //    printf("
fin
");
 75     if(id>1 && id<cnt){//FFT不统计边缘 
 76         memset(a,0,sizeof a);
 77         memset(b,0,sizeof b);
 78         for(i=1;i<B[id].L;i++){
 79             a[w[i]].x+=1;
 80         }
 81         for(i=B[id].R+1;i<=n;i++){
 82             b[w[i]].x+=1;
 83         }
 84         FFT(a,1);FFT(b,1);
 85         for(i=0;i<N;i++){
 86             a[i]=a[i]*b[i];
 87         }
 88         FFT(a,-1);
 89         for(i=B[id].L;i<=B[id].R;i++){
 90             ans+=(LL)(a[w[i]<<1].x+0.3);
 91         }
 92     }
 93     for(i=B[id].L;i<=B[id].R;i++)
 94         ++Lc[w[i]];
 95     return;
 96 }
 97 int main(){
 98 //    freopen("in.txt","r",stdin);
 99     int i,j;
100     n=read();
101     for(i=1;i<=n;i++)w[i]=read();
102     block=2000;
103     cnt=0;
104     
105     while(cnt*block<n){
106         B[++cnt].L=B[cnt-1].R+1;
107         B[cnt].R=cnt*block;
108         B[cnt].id=cnt;
109     }
110     B[cnt].R=min(B[cnt].R,n);
111     int m=60000;
112     for(N=1,l=0;N<=m;N<<=1)l++;
113     R_init(N);
114     for(i=n;i;i--)Rc[w[i]]++;
115     for(i=1;i<=cnt;i++)
116         solve(i);
117     //
118     printf("%lld
",ans);
119     return 0;
120 }