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

9

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

 

/*by SilverN*/
#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
#include<vector>
#define LL long long
using namespace std;
const double pi=acos(-1.0);
const int mxn=120010;
int read(){
    int x=0,f=1;char ch=getchar();
    while(ch<'0' || ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0' && ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
//
struct com{
    double x,y;
    com operator + (com b){return (com){x+b.x,y+b.y};}
    com operator - (com b){return (com){x-b.x,y-b.y};}
    com operator * (com b){return (com){x*b.x-y*b.y,x*b.y+y*b.x};}
}a[mxn],b[mxn],c[mxn];
int rev[mxn];
int N,l;
void R_init(int n){
    for(int i=0;i<n;i++)rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
    return;
}
void FFT(com *a,int flag){
    for(int i=0;i<N;i++)if(rev[i]>i)swap(a[rev[i]],a[i]);
    for(int i=1;i<N;i<<=1){
        com wn=(com){cos(pi/i),flag*sin(pi/i)};
        for(int j=0;j<N;j+=(i<<1)){
            com w=(com){1,0};
            for(int k=0;k<i;k++,w=w*wn){
                com x=a[j+k],y=w*a[i+j+k];
                a[j+k]=x+y;
                a[i+j+k]=x-y;
            }
        }
    }
    if(flag==-1){for(int i=0;i<N;i++)a[i].x/=N;}
    return;
}
//FFT
struct node{
    int L,R;
    int id;
}B[mxn];
int block,cnt;
//分块 
int w[mxn];
LL ans=0;
int Lc[mxn],Rc[mxn];
int n;
void solve(int id){
//    printf("solve:%d\n",id);
//    printf("%lld\n",ans);
    int i,j,x;
    for(i=B[id].L;i<=B[id].R;i++){//统计块内 
        --Rc[w[i]];
        for(j=B[id].L;j<i;j++){
            x=w[i]*2-w[j];
            if(x>=0)ans+=Rc[x];
        }
        for(j=i+1;j<=B[id].R;j++){
            x=w[i]*2-w[j];
            if(x>=0)ans+=Lc[x];
        }
    }
//    return;
//    printf("\nfin\n");
    if(id>1 && id<cnt){//FFT不统计边缘 
        memset(a,0,sizeof a);
        memset(b,0,sizeof b);
        for(i=1;i<B[id].L;i++){
            a[w[i]].x+=1;
        }
        for(i=B[id].R+1;i<=n;i++){
            b[w[i]].x+=1;
        }
        FFT(a,1);FFT(b,1);
        for(i=0;i<N;i++){
            a[i]=a[i]*b[i];
        }
        FFT(a,-1);
        for(i=B[id].L;i<=B[id].R;i++){
            ans+=(LL)(a[w[i]<<1].x+0.3);
        }
    }
    for(i=B[id].L;i<=B[id].R;i++)
        ++Lc[w[i]];
    return;
}
int main(){
//    freopen("in.txt","r",stdin);
    int i,j;
    n=read();
    for(i=1;i<=n;i++)w[i]=read();
    block=2000;
    cnt=0;
    
    while(cnt*block<n){
        B[++cnt].L=B[cnt-1].R+1;
        B[cnt].R=cnt*block;
        B[cnt].id=cnt;
    }
    B[cnt].R=min(B[cnt].R,n);
    int m=60000;
    for(N=1,l=0;N<=m;N<<=1)l++;
    R_init(N);
    for(i=n;i;i--)Rc[w[i]]++;
    for(i=1;i<=cnt;i++)
        solve(i);
    //
    printf("%lld\n",ans);
    return 0;
}