bzoj3509[CodeChef]COUNTARI

题意

给出一个数组a[1...n],问有多少三元组(i,j,k)满足i<j<k且a[i]+a[k]==2*a[j]

分析

考虑暴力,枚举j,那么需要在两侧各找出一个i,k,使得(a[i]+a[k]==2*a[j]),那么我们把j左边的所有数字找出来,右边的所有数字找出来,找出有哪些数字对的和为(2*a[j]), 这是一个卷积形式所以可以FFT,第(2*a[j])项的系数就是答案.那么这样的复杂度是(O(n^2logn))的,无法承受.

仔细考虑这个过程,其实有很多冗余的计算.比如说枚举到j=100和j=101的时候,卷积的结果变化很小.但是并不能方便地从j=100的卷积结果快速推出j=101的卷积结果.从另一个角度思考,j=100和j=101的时候都需要考虑(a[1...99])和(a[102...n])这两段数字之间的卷积,但这两部分结果我们其实计算了两次.如何减少这个计算的次数?合理的方法是分块.

将a[]数组分成block块,我们分别考虑i,j,k在3个不同的块的情况和i,j,k中至少有两个在同一个块的情况.

i,j,k在3个不同的块时

我们枚举j所在的块,那么i在左侧的j-1个块中,k在右侧的block-j个块中,我们将左侧的j-1个块中每个数字的出现次数统计出来,右侧block-j个块中每个数字的出现次数统计出来(实现的时候只需每扫过一个块就用这个块内的数字更新左侧和右侧的统计结果),做一遍fft即可.共需block-2次fft(j不可能出现在第1或第block个块).

如果块大小取(/{sqrt}(n)),这一部分的复杂度为(O(nlog(n)sqrt(n)))

i,j,k至少有两个在同一个块x之内时

我们枚举x是哪一个块,然后枚举其中的两个元素i,j(不妨令i<j,且假定k不在i和j之间),那么考虑k和i,j是否在同一个块内.不在同一个块内,如果k在前面那么(a[k]=a[i]*2-a[j]),如果k在后面那么(a[k]=a[j]*2-a[i]),我们边扫边维护所有前面的块中数字的出现次数和后面的块中数字出现的次数.对于k和i,j在同一个块内的情况,我们可以在枚举i,j的时候适当维护同一个块中j后面的数字的出现次数(可以合并到后面的块中的数字出现次数,一起维护),就可以一起更新答案.如果块大小取(/{sqrt}(n)),这一部分的复杂度为(O(nsqrt(n)))

实际上,由于fft的常数较大,块大小应当取得大一些以减少fft次数

一开始我的写法比较丑,间接统计答案,需要2*blocks次fft,卡了半天常还是没过虽然复杂度一样就不放上来了,贴一个能过的

如果仔细观察里面还有各种卡常数的痕迹

#include<cstdio>
#include<algorithm>
#include<cmath>
#include<ctime>
#include<cctype>
void read(int &x){
  char ch;while(ch=getchar(),!isdigit(ch));x=ch-'0';
  while(ch=getchar(),isdigit(ch))x=x*10+ch-'0';
}
using namespace std;
typedef long long ll;
const double pi=acos(-1);
const int maxn=100050;
int N=65536;
struct comp{
  double x,y;
  comp(){}
  comp(double a,double b){x=a;y=b;}
  comp operator +(const comp &a)const{return comp(x+a.x,y+a.y);}
  comp operator -(const comp &a)const{return comp(x-a.x,y-a.y);}
  comp operator *(const comp &a)const{return comp(x*a.x-y*a.y,a.x*y+a.y*x);}
  void operator +=(const comp &a){x+=a.x;y+=a.y;}
};
comp wn[2][17][maxn];
//double sum=0;
void fft(comp *a,int N,int sign){
  // double t1=clock();
  for(int i=1,j=0,k=N;i<N;++i,k=N){
    do j^=(k>>=1);while(j<k);
    if(i<j)swap(a[i],a[j]);
  }
  for(int j=2,i=0;j<=N;j<<=1,++i){
    int m=j>>1;comp *W=wn[sign][i];
    for(comp* p=a;p!=a+N;p+=j){
      for(int i=0;i<m;++i){
	comp t=p[m+i]*W[i];
	p[m+i]=p[i]-t;p[i]+=t;
      }
    }
  }
  if(sign==1)for(int i=0;i<N;++i)a[i].x/=N;
  //double t2=clock();sum+=t2-t1;
}
void mul(ll *a,ll* b,ll *c){
  static comp A[maxn],B[maxn];
  for(int i=0;i<N;++i)A[i]=comp(a[i],0),B[i]=comp(b[i],0);
  fft(A,N,0);fft(B,N,0);
  for(int i=0;i<N;++i)A[i]=A[i]*B[i];
  fft(A,N,1);
  for(int i=0;i<N;++i)c[i]=(ll)floor(A[i].x+0.5);
}
int sz=4000;
int L[maxn],R[maxn],belong[maxn];
ll ans=0;
ll suc[maxn],pre[maxn];
int a[maxn];
ll f[maxn];
void brute(int x){
  if(x>1)for(int i=L[x-1];i<=R[x-1];++i)pre[a[i]]++;
  for(int i=L[x];i<=R[x];++i){
    suc[a[i]]--;
    for(int j=L[x];j<i;++j){
      if(0<=2*a[i]-a[j]){
	ans+=suc[2*a[i]-a[j]];
      }
      if(0<=2*a[j]-a[i])ans+=pre[2*a[j]-a[i]];
    }
  }
}
int main(){
  int n;scanf("%d",&n);int Max=0;
  for(int i=1;i<=n;++i){
    scanf("%d",&a[i]);if(a[i]>Max)Max=a[i];
  }
  while(Max*4<=N&&N>=100)N>>=1;
  // double t3=clock();
  for(int j=2,i=0;j<=N;j<<=1,++i){
    comp tmp=comp(cos(2*pi/j),sin(2*pi/j));
    wn[0][i][0]=comp(1,0);
    if(i>=1){
      for(int k=1;k<j;++k){
	if(k&1)wn[0][i][k]=wn[0][i][k-1]*tmp;
	else wn[0][i][k]=wn[0][i-1][k>>1];
      }
    }else{
      for(comp* pt=(&wn[0][i][1]),*ed=(&wn[0][i][j]);pt!=ed;++pt){
	(*pt)=(*(pt-1))*tmp;
      }
    }
    tmp=comp(cos(2*pi/j),-sin(2*pi/j));
    wn[1][i][0]=comp(1,0);
    if(i>=1){
      for(int k=1;k<j;++k){
	if(k&1)wn[1][i][k]=wn[1][i][k-1]*tmp;
	else wn[1][i][k]=wn[1][i-1][k>>1];
      }
    }else{
      for(comp* pt=(&wn[1][i][1]),*ed=(&wn[1][i][j]);pt!=ed;++pt){
	(*pt)=(*(pt-1))*tmp;
      }
    }
  }
  int tot=(n+sz-1)/sz;
  for(int i=1;i<tot;++i){
    L[i]=(i-1)*sz+1;R[i]=i*sz;
  }
  L[tot]=(tot-1)*sz+1;R[tot]=n;
  for(int i=1;i<=tot;++i){
    for(int j=L[i];j<=R[i];++j){
      belong[j]=i;
    }
  }
  for(int i=1;i<=n;++i)suc[a[i]]++;
  for(int i=1;i<=tot;++i)brute(i);
  for(int i=L[2];i<=n;++i)suc[a[i]]++;
  for(int i=0;i<N;++i)pre[i]=0;
  for(int i=2;i<tot;++i){
    for(int j=L[i-1];j<=R[i-1];++j)pre[a[j]]++;
    for(int j=L[i];j<=R[i];++j)suc[a[j]]--;
    mul(pre,suc,f);
    for(int j=L[i];j<=R[i];++j){
      ans+=f[a[j]*2];
    }
  }
  printf("%lld
",ans);
  return 0;
}