[BZOJ 3509] [CodeChef] COUNTARI (FFT+分块)
题面
给出一个长度为n的数组,问有多少三元组\((i,j,k)\)满足\(i<j<k,a_j-a_i=a_k-a_j\)
\(n \leq 10^5, a_i \leq 30000\)
分析
记\(m=\max(a_i)\)
先做一点变形\(a_i+a_k=2a_j\).那么枚举j,把j左边的和右边的所有数字找出来,找出有哪些数字对的和为\(2a_j\)。
维护当前块左边每个数的出现次数cntl,，和右边块的出现次数cntr。
那么数字和为\(2a_j\)的方案数就是\(\sum_{i=0}^{2a_j} cntl(i)cntr(2a_j-i)\)
这是一个卷积的形式,可以FFT.令:
$A(x)=\sum_{i=0}^{m} cntl(i) x^i $
$B(x)=\sum_{i=0}^{m} cntr(i) x^i $
\(A*B\)的第\(2a_j\)项的系数就是答案.但是做n次FFT的复杂度是\(O(n^2 \log n)\),显然是跑不过去的.
我们要平衡一下复杂度，容易想到分块。设块大小为T
i,j,k均不在同一个块的情况
我们枚举j在第几个块,从左到右扫描第2~n-1个块,每块把左边和右边FFT卷积。枚举块内的每个数x,和为\(2x\)的方案数就是卷积结果里第\(2x\)项的系数.
时间复杂度\(O(\frac{n}{B} \times m \log m)\)
i,j,k至少有2个在同一个块的情况
直接暴力从小到大在块内枚举\(k,j(j<k)\).然后讨论j在k前还是在k后。同样维护2个数组cntl表示当前块前面的数的出现次数，cntr表示k后面的数的个数。
如果i在j前,那么答案就是\(cntl[a_j*2-a_k]\)
如果i在k后，那么答案就是\(cntr[a_k*2-a_j]\)
注意到cntl不包含当前块，而cntr包含当前块，这样\(i,j,k\)在同一个块的情况只会被算1次
时间复杂度\(O(\frac{n}{B} \times B \times B)=O(nB)\)
总时间复杂度\(O(nB+\frac{n}{B}m\log m)\)
根据基本不等式,\(B\)取\(\sqrt{m \log m}\)时最优.总时间复杂\(O(n\sqrt{m \log m})\)
代码
#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#define maxn 200000
using namespace std;
typedef long double db;
typedef long long ll;
const db pi=acos(-1.0);
struct com{
double real;
double imag;
com(){
}
com(double _real,double _imag){
real=_real;
imag=_imag;
}
com(double x){
real=x;
imag=0;
}
void operator = (const com x){
this->real=x.real;
this->imag=x.imag;
}
void operator = (const double x){
this->real=x;
this->imag=0;
}
friend com operator + (com p,com q){
return com(p.real+q.real,p.imag+q.imag);
}
friend com operator + (com p,double q){
return com(p.real+q,p.imag);
}
void operator += (com q){
*this=*this+q;
}
void operator += (double q){
*this=*this+q;
}
friend com operator - (com p,com q){
return com(p.real-q.real,p.imag-q.imag);
}
friend com operator - (com p,double q){
return com(p.real-q,p.imag);
}
void operator -= (com q){
*this=*this-q;
}
void operator -= (double q){
*this=*this-q;
}
friend com operator * (com p,com q){
return com(p.real*q.real-p.imag*q.imag,p.real*q.imag+p.imag*q.real);
}
friend com operator * (com p,double q){
return com(p.real*q,p.imag*q);
}
void operator *= (com q){
*this=(*this)*q;
}
void operator *= (double q){
*this=(*this)*q;
}
friend com operator / (com p,double q){
return com(p.real/q,p.imag/q);
}
void operator /= (double q){
*this=(*this)/q;
}
void print(){
printf("%lf + %lf i ",real,imag);
}
};
void fft(com *x,int n,int type){
static int rev[maxn+5];
int dn=1,k=0;
while(dn<n){
dn*=2;
k++;
}
for(int i=0;i<n;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(k-1));
for(int i=0;i<n;i++) if(i<rev[i]) swap(x[i],x[rev[i]]);
for(int len=1;len<n;len*=2){
int sz=len*2;
com wn1=com(cos(2*pi/sz),sin(2*pi/sz)*type);
for(int l=0;l<n;l+=sz){
int r=l+len-1;
com wnk=1;
for(int i=l;i<=r;i++){
com tmp=x[i+len];
x[i+len]=x[i]-wnk*tmp;
x[i]=x[i]+wnk*tmp;
wnk*=wn1;
}
}
}
if(type==-1) for(int i=0;i<n;i++) x[i]/=n;
}
void mul(com *a,com *b,com *ans,int n){
fft(a,n,1);
if(a!=b) fft(b,n,1);
for(int i=0;i<n;i++) ans[i]=a[i]*b[i];
fft(ans,n,-1);
}
int n;
int bsz,bcnt;
int a[maxn+5];
int bel[maxn+5];
int lb[maxn+5],rb[maxn+5];//预处理每个数的左右边界，卡常
void ini_block(){
bsz=sqrt(n*log(n)/log(2));
bcnt=1;
for(int i=1;i<=n;i++){
bel[i]=bcnt;
if(i%bsz==0) bcnt++;
}
for(int i=1;i<=bcnt;i++){
lb[i]=bsz*(i-1)+1;
rb[i]=min(i*bsz,n);
}
}
int cntl[maxn+5],cntr[maxn+5];
com fl[maxn+5],fr[maxn+5],res[maxn+5];
int main(){
scanf("%d",&n);
ini_block();
int maxv=0;
for(int i=1;i<=n;i++){
scanf("%d",&a[i]);
maxv=max(maxv,a[i]);
}
//i,j,k有>=2个在同一块内，直接暴力O(n/B*B^2)=O(Bn)
ll ans=0;
for(int i=n;i>=1;i--) cntr[a[i]]++;
for(int id=1;id<=bcnt;id++){
//cntl记录的是当前块前面的数的出现情况
//cntr记录的是i后面的数的出现情况
//注意i,j,k均在一个块内的情况只能被算1次
for(int k=lb[id];k<=rb[id];k++){
cntr[a[k]]--;
for(int j=lb[id];j<k;j++){
int ai=a[j]*2-a[k];
if(ai>=0&&ai<=maxn) ans+=cntl[ai];
ai=a[k]*2-a[j];
if(ai>=0&&ai<=maxn) ans+=cntr[ai];//这样只有这个时候会算到i,j,k在同一个块里的情况
}
}
for(int j=lb[id];j<=rb[id];j++) cntl[a[j]]++;
}
memset(cntl,0,sizeof(cntl));
memset(cntr,0,sizeof(cntr));
for(int i=n;i>=lb[2];i--) cntr[a[i]]++;
for(int i=1;i<=rb[1];i++) cntl[a[i]]++;
for(int id=2;id<bcnt;id++){
for(int j=lb[id];j<=rb[id];j++) cntr[a[j]]--;
for(int j=0;j<=maxv;j++){
fl[j]=cntl[j];
fr[j]=cntr[j];
}
int dn=1;
while(dn<=maxv*2) dn*=2;
mul(fl,fr,res,dn);
for(int j=lb[id];j<=rb[id];j++){
ans+=(ll)(res[a[j]*2].real+0.5);
}
for(int j=0;j<=dn;j++){
fl[j]=0;
fr[j]=0;
res[j]=0;
}
for(int j=lb[id];j<=rb[id];j++) cntl[a[j]]++;
}
printf("%lld\n",ans);
}