【BZOJ3509】[CodeChef] COUNTARI（分块好题）
大致题意： 给定一个序列，问有多少对\(i,j,k\)满足\(i<j<k\)且\(a_k-a_j=a_j-a_i\)。
从暴力到分块
看到\(a_k-a_j=a_j-a_i\)，我们容易想到把它移项得到\(a_i+a_k=2a_j\)。
于是有个很暴力的想法：枚举\(j\)，然后统计\(fr_x,bk_x\)分别表示\(j\)前后\(a_i=x\)的\(i\)的个数，则每次对\(fr\)和\(bk\)做个\(FFT\)，其中第\(2a_j\)位的系数就是能与\(j\)匹配的\(i,k\)对数。
这个做法的正确性应该是显然的，而且这个做法复杂度很大会T飞也是显然的。
考虑能不能对这个算法进行优化，然后发现相邻若干次\(FFT\)其实\(fr,bk\)变化很小，有很多计算都是白白浪费的。
因此，我们祭出分块大法，假设把原序列分成若干大小为\(S\)的块，然后分情况讨论。
情况\(1\)：三点分属于不同块中
显然此时我们可以枚举\(j\)在哪一个块（假设为第\(t\)个块），然后\(i,k\)就分别只能在第\(1\sim t-1\)个块和第\(t+1\sim n\)个块中。
于是我们再拿出前面的\(fr_x,bk_x\)，分别表示第\(1\sim t-1\)个块和第\(t+1\sim n\)个块中\(a_i=x\)的\(i\)的个数。
按照暴力的想法我们对\(fr_x,bk_x\)做个\(FFT\)，然后枚举块中的\(j\)计算答案。
时间复杂度：\(O(\frac nS\times nlogn)\)。
情况\(2\)：三点中存在属于同一块中的点
假设枚举\(p,q(p<q)\)两点属于同一块（另一点\(w\)可以不在同一块，也可以在同一块，但不能在\(p,q\)中间）。
我们同样统计\(fr_x,bk_x\)（我想应该不用再重复一次它们的定义了。。。），于是又分三种情况：
- \(w\)在之前的块：则\(p\)对应\(j\)，给答案加上\(fr_{2\times a_p-a_q}\)。
- \(w\)在之后的块：则\(q\)对应\(j\)，给答案加上\(bk_{2\times a_q-a_p}\)。
- \(w\)在同一块：显然我们只需考虑\(w<p\)的情况（若\(w>q\)，则在枚举到\(q,w\)时会重复计算贡献），然后我们发现实际上只要把\(fr_x\)改为统计\(1\sim p-1\)的所有数中\(a_i=x\)的\(i\)的个数，就可以把这种情况和\(w\)在之前的块的情况一并计算了，而这个统计显然是可以在枚举\(p\)的过程中同时维护的。
时间复杂度：\(O(nS)\)。
关于块的大小\(S\)
我们发现这道题中两种情况下复杂度是不等的。
我们令\(\frac nS\times nlogn=nS\)，然后就得到 \(S=\sqrt{nlogn}\)，实测了几种块的大小似乎这也的确是最优的。
代码
#pragma GCC optimize(2)
#pragma GCC optimize("inline")
#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define N 100000
#define V 30000
#define DB double
using namespace std;
int n,sz,a[N+5],fr[2*V+5],bk[2*V+5];long long res[2*V+5];
class Poly//FFT模板
{
private:
int P,L,R[V<<2];DB Pi;
struct node
{
DB x,y;I node(Con DB& a=0,Con DB& b=0):x(a),y(b){}
I node operator + (Con node& o) Con {return node(x+o.x,y+o.y);}
I node operator - (Con node& o) Con {return node(x-o.x,y-o.y);}
I node operator * (Con node& o) Con {return node(x*o.x-y*o.y,x*o.y+y*o.x);}
}A[V<<2],B[V<<2];
I void T(node *s,CI op)
{
RI i,j,k;node U,S,x,y;for(i=0;i^P;++i) i<R[i]&&(x=s[i],s[i]=s[R[i]],s[R[i]]=x,0);
for(i=1;i^P;i<<=1) for(U=node(cos(Pi/i),op*sin(Pi/i)),j=0;j^P;j+=i<<1)
for(S=node(1,0),k=0;k^i;++k,S=S*U) s[j+k]=(x=s[j+k])+(y=S*s[i+j+k]),s[i+j+k]=x-y;
}
public:
I Poly()
{
Pi=acos(-1),P=1,L=0;W(P<=(V<<1)) P<<=1,++L;
for(RI i=0;i^P;++i) R[i]=(R[i>>1]>>1)|((i&1)<<L-1);
}
I void FFT(int *x,int *y,long long *res)
{
RI i;for(i=0;i^P;++i) A[i]=B[i]=node();
for(i=0;i<=V;++i) A[i]=x[i],B[i]=y[i];
for(T(A,1),T(B,1),i=0;i^P;++i) A[i]=A[i]*B[i];
for(T(A,-1),i=0;i<=(V<<1);++i) res[i]=A[i].x/P+0.5;
}
}P;
int main()
{
RI i,j,k,lim;long long ans=0;scanf("%d",&n),sz=sqrt(n*(log(n)/log(2)));
for(i=1;i<=n;++i) scanf("%d",a+i);
for(i=1;i<=sz;++i) ++fr[a[i]];for(i=sz+1;i<=n;++i) ++bk[a[i]];//统计fr,bk（方便起见一开始把第2个块统计在bk）
for(i=2;i*sz<n;++i)//情况1（注意无需枚举首尾两个块）
{
for(j=(i-1)*sz+1;j<=i*sz;++j) --bk[a[j]];P.FFT(fr,bk,res);//删去当前块的贡献，然后FFT
for(j=(i-1)*sz+1;j<=i*sz;++j) ++fr[a[j]],ans+=res[2*a[j]];//加上当前块的贡献，同时统计答案
}
for(i=0;i<=V;++i) fr[i]=bk[i]=0;for(i=1;i<=n;++i) ++bk[a[i]];//清空，重新统计bk
for(i=1;i<=(n-1)/sz+1;++i)//情况2
{
for(lim=min(i*sz,n),j=(i-1)*sz+1;j<=lim;++j) --bk[a[j]];//删去当前块的贡献
for(j=(i-1)*sz+1;j<=lim;++j)//枚举一对点
{
for(k=j+1;k<=lim;++k) 2*a[j]>=a[k]&&
(ans+=fr[2*a[j]-a[k]]),2*a[k]>=a[j]&&(ans+=bk[2*a[k]-a[j]]);//统计答案
++fr[a[j]];//加上当前点贡献
}
}return printf("%lld",ans),0;
}
