[BZOJ3509][CodeChef][FFT][分块]COUNTARI

分块FFT练习题

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <cstring>
#include <string>
#include <cmath>
#define N 200010

using namespace std;

typedef long long ll;

struct E{
  double real,imag;
  E(double r=0,double i=0):real(r),imag(i){}
  inline friend E operator +(E a,E b){
    return E(a.real+b.real,a.imag+b.imag);
  }
  inline friend E operator -(E a,E b){
    return E(a.real-b.real,a.imag-b.imag);
  }
  inline friend E operator *(E a,E b){
    return E(a.real*b.real-a.imag*b.imag,a.real*b.imag+a.imag*b.real);
  }
  inline friend E operator /(E a,double b){
    return E(a.real/b,a.imag/b);
  }
};

int A[N],pos[N],L[N],R[N],st[N],ed[N],rev[600000],pre[N];
ll c[600000];
E a[600000],b[600000];
int n,block,g,m,M,v,len;
ll Ans;

inline void reaD(int &x){
  char c=getchar();x=0;
  for(;c>'9'||c<'0';c=getchar());for(;c>='0'&&c<='9';x=x*10+c-'0',c=getchar());
}

inline void swap(E &a,E &b){
  E c=a; a=b; b=c;
}

inline void FFT(E *a,int w){
  for(int i=0;i<m;++i) if(rev[i]>i) swap(a[i],a[rev[i]]);
  for(int i=1;i<m;i<<=1){
    E wn(cos(M_PI/i),w*sin(M_PI/i));
    for(int j=0;j<m;j+=(i<<1)){
      E w(1,0);
      for(int k=0;k<i;k++,w=w*wn){
    E x=a[j+k],y=w*a[j+k+i];
    a[j+k]=x+y; a[j+k+i]=x-y;
      }
    }
  }
  if(w==-1) for(int i=0;i<m;i++) a[i]=a[i]/m;
}

inline void _fft(){
  memset(a,0,sizeof(a));
  memset(b,0,sizeof(b));
  for(int i=0;i<=v;++i)
    a[i]=L[i],b[i]=R[i];
  FFT(a,1); FFT(b,1);
  for(int i=0;i<m;++i) a[i]=a[i]*b[i];
  FFT(a,-1);
  for(int i=0;i<=M;i++) c[i]=(ll)(a[i].real+0.1);
}

inline void init(){
  for(m=1,M=v<<1;m<=M;m<<=1) len++;
  for(int i=0;i<m;++i) rev[i]=(rev[i>>1]>>1)|((i&1)<<(len-1));
}

inline void solve(){
  init();
  for(int i=1;i<=g;++i){
    for(int j=st[i-1];i-1&&j<=ed[i-1];++j) L[A[j]]++;
    for(int j=st[i];j<=ed[i];++j) R[A[j]]--;
    if(i>1&&i<g) _fft();
    memset(pre,0,sizeof(pre));
    for(int j=st[i];j<=ed[i];++j){
      int t=A[j]+A[j];
      for(int k=st[i];k<j;++k)
    if(t-A[k]>0&&t-A[k]<=v) Ans+=R[t-A[k]];
      for(int k=j+1;k<=ed[i];++k)
    if(t-A[k]>0&&t-A[k]<=v) Ans+=L[t-A[k]]+pre[t-A[k]];
      if(i>1&&i<g) Ans+=c[t]; pre[A[j]]++;
    }
  }
}

int main(){
  reaD(n); block=2000;
  for(int i=1;i<=n;++i){
    reaD(A[i]);
    pos[i]=(i-1)/block+1;
    R[A[i]]++;
    v=max(v,A[i]);
  }
  g=pos[n]; 
  for(int i=1;i<=g;++i)
    st[i]=(i-1)*block+1,ed[i]=i*block;
  ed[g]=min(ed[g],n);
  solve();
  printf("%lld\n",Ans);
}