3509: [CodeChef] COUNTARI

Time Limit: 40 Sec Memory Limit: 128 MB
Submit: 774 Solved: 216
[Submit][Status][Discuss]

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

[Submit][Status][Discuss]

------------------------

写这个题解主要想吐槽一下bzoj的评测机

不是因为他慢而是他有毒

我第一次拿fft做有一点不裸的题就遇到这样的情况 T T

从上往下数第二个是我从网上搞过来的标程

我tm就不服了向bzoj要了数据自己测

发现

才差了0.01s 没开o2 而且电脑比较破但也就40s

我怀疑帮我评测的机器是假的

算了不管他就当自己a了~

至于怎么做

先考虑暴力fft

每个点把他左边和他右边的生成函数乘一下

取结果中的(a[i]*2)那一项的系数

复杂度为O(n^2logn)

然后看看感觉每个数都要fft最后还只取一项岂不是很奢侈

然后就分块每一块做fft

比如当前块 [L,R]

把 [1,L)和(R,n)的生成函数乘一下

块内的每个数字都取需要的那一项就好了

最后每个块内自己暴力搞一下就好了

(说了半天还不是看了别人题解才会做)

代码：(没什么用反正交上去会T 欢迎大家来试一试)

#include<cstdio>
#include<cmath>
#include<iostream>
#include<cstring>
#define For(i,x,y) for(int i=x;i<=y;++i)
using namespace std;
const int N = 4e5 + 5;
const double pi = 3.14159265358979323846264338327950;
const double eps = 1e-2;
#define min(a,b) ((a)<(b)?(a):(b)) 
struct Complex{
    double r,i;
    Complex operator + (Complex b){return (Complex){r+b.r,i+b.i};}
    Complex operator - (Complex b){return (Complex){r-b.r,i-b.i};}
    Complex operator * (Complex b){return (Complex){r*b.r-i*b.i,r*b.i+i*b.r};}
    Complex operator ^ (int f){return (Complex){r,i*f};}
}w[N],a[N],b[N],c[N],tmp[N];
int p[N];
inline void init_w(int n){For(i,0,n-1){w[i]=(Complex){cos(2*pi*i/n),sin(2*pi*i/n)};}}
inline void fft(int n,Complex*buf,int beg,int step,int f){
    if(n==1) return;
    int m=n>>1;
    fft(m,buf,beg,step<<1,f);
    fft(m,buf,beg+step,step<<1,f);
    For(i,0,m-1){
        int pos=i*step*2;
        tmp[i]=buf[beg+pos]+(w[i*step]^f)*(buf[beg+pos+step]);
        tmp[i+m]=buf[beg+pos]-(w[i*step]^f)*(buf[beg+pos+step]);
    }
    For(i,0,n-1) buf[beg+step*i] = tmp[i];
}
int n;
int Bg;
int block_size;
long long ans=0;
void solve_out(){
    for(int l=1;l<=n;l+=block_size){
        int m=-1;        int r=min(n,l+block_size-1);
        For(i,0,Bg) a[i]=b[i]=(Complex){0,0};
        For(i,1,l-1) a[p[i]].r+=1,m=max(p[i],m);
        For(i,r+1,n) b[p[i]].r+=1,m=max(p[i],m);
        int len=1;for(;len<=m+m;len<<=1);
//        len<<=1;
        init_w(len);
        fft(len,a,0,1,1);
        fft(len,b,0,1,1);
        For(i,0,len) a[i]=a[i]*b[i];
        fft(len,a,0,1,-1);
        For(i,l,r){
            ans+=(long long)((a[p[i]<<1].r+0.01)/len);
        }
    }
}
int L[N],R[N];
void solve_in(){
    For(i,1,n) R[p[i]]++;
    for(int l=1;l<=n;l+=block_size){
        int r=min(n,l+block_size-1);
        For(i,l,r) R[p[i]]--;
        For(i,l,r){
            For(j,i+1,r){
                if(2*p[i]-p[j]>=0)ans+=L[2*p[i]-p[j]];
                if(2*p[j]-p[i]>=0)ans+=R[2*p[j]-p[i]];
            }
            L[p[i]]++;
        }
    }
}
int main()
{
//    freopen("3509.in","r",stdin);
//    freopen("3509.out","w",stdout);
    scanf("%d",&n);
    For(i,1,n) scanf("%d",&p[i]),Bg=max(Bg,p[i]);
    int rwy233=Bg;Bg=1; while(Bg<rwy233) Bg<<=1; 
    Bg<<=1;
    block_size=(int)sqrt(n)*16;
    if(rwy233<=1000) block_size = 2000;
//    block_size=3;
    if(block_size<1) block_size=1;
    if(block_size>n) block_size=n;
    solve_out();
    solve_in();
    printf("%lld",ans);
    return 0;
}