[Codechef November Challenge 2012] Arithmetic Progressions

 1 #include<bits/stdc++.h>
 2 using namespace std;
 3 #define LL long long
 4 #define N 205005
 5 #define INF 1e9
 6 #define Bl 70
 7 #define LIM 60000
 8 const double PI=acos(-1);
 9 
10 inline LL read(){
11        LL x=0,f=1; char a=getchar();
12        while(a<‘0‘ || a>‘9‘) {if(a==‘-‘) f=-1; a=getchar();}
13        while(a>=‘0‘ && a<=‘9‘) x=x*10+a-‘0‘,a=getchar();
14        return x*f;
15 }
16 
17 namespace FFT{
18     int rev[N];
19     
20     struct vec{
21         double r,i;
22         vec operator * (const vec& w){return (vec){r*w.r-i*w.i,i*w.r+r*w.i};}
23         vec operator + (const vec& w){return (vec){r+w.r,i+w.i};}
24         vec operator - (const vec& w){return (vec){r-w.r,i-w.i};}
25     }A[N],B[N];
26     
27     inline void fft(vec* x,int len,int f){
28         for(int i=1;i<=len;i++) if(i<rev[i]) swap(x[i],x[rev[i]]);
29         for(int lnow=2;lnow<=len;lnow<<=1){
30             vec w,w0=(vec){cos(2.0*PI/lnow*f),sin(2.0*PI/lnow*f)},t1,t2;
31             for(int i=0;i<len;i+=lnow){
32                 w=(vec){1,0};
33                 for(int j=i;j<i+lnow/2;j++){
34                     t1=x[j]; t2=w*x[j+lnow/2];
35                     x[j]=t1+t2; x[j+lnow/2]=t1-t2;
36                     w=w*w0;
37                 }
38             }
39         }
40     }
41     
42     inline void work(int a[],int b[],int l1,int l2,LL s[]){
43         int len,t;
44         for(len=1,t=0;len<=(l1+l2+1);len<<=1,t++); t=1<<(t-1);
45         for(int i=0;i<=len;i++) rev[i]=rev[i>>1]>>1|(i&1?t:0);
46         for(int i=0;i<=len;i++) B[i]=A[i]=(vec){0,0};
47         for(int i=0;i<=l1;i++) A[i].r=a[i];
48         for(int i=0;i<=l2;i++) B[i].r=b[i];
49         fft(A,len,1); fft(B,len,1);
50         for(int i=0;i<=len;i++) A[i]=A[i]*B[i];
51         fft(A,len,-1);
52         for(int i=0;i<=l1+l2;i++)
53         s[i]=(LL)(1.0*A[i].r/len+0.5);
54     }
55     
56 }
57 
58 int n,block_size,block_num;
59 int bel[N],l[Bl+5],r[Bl+5],a[N];
60 LL tot,ans[2*LIM+6000];
61 int lsum[LIM+5],rsum[LIM+5],cnt[2*LIM+5];
62 
63 inline void brutal_force(int x){
64     for(int i=l[x];i<=r[x];i++) rsum[a[i]]--;
65     memset(ans,0,sizeof(ans));
66     FFT::work(lsum,rsum,30000,30000,ans);
67     for(int i=l[x];i<=r[x];i++){
68         tot+=ans[2*a[i]];
69         for(int j=l[x];j<i;j++)
70         if(2*a[i]-a[j]>0) tot+=rsum[2*a[i]-a[j]];
71         for(int j=i+1;j<=r[x];j++)
72         if(2*a[i]-a[j]>0) tot+=lsum[2*a[i]-a[j]];
73     }
74     for(int i=l[x];i<=r[x];i++) lsum[a[i]]++;
75     memset(cnt,0,sizeof(cnt));
76     for(int i=l[x];i<=r[x];i++){
77         tot+=cnt[a[i]];
78         for(int j=l[x];j<i;j++)
79         if(2*a[i]-a[j]>0) cnt[2*a[i]-a[j]]++;
80     }
81 }
82 
83 int main(){
84     n=read(); block_size=1500;
85     block_num=(n-1)/block_size+1;
86     for(int i=1;i<=n;i++) a[i]=read(),bel[i]=(i-1)/block_size+1;
87     for(int i=1;i<=block_num;i++) l[i]=(i-1)*block_size+1,r[i]=min(n,i*block_size);
88     for(int i=1;i<=n;i++) rsum[a[i]]++;
89     for(int i=1;i<=block_num;i++) brutal_force(i);
90     printf("%lld\n",tot);
91     return 0;
92 }