大致题意: 给定一个序列,问有多少对(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)就分别只能在第(1sim t-1)个块和第(t+1sim n)个块中。
于是我们再拿出前面的(fr_x,bk_x),分别表示第(1sim t-1)个块和第(t+1sim n)个块中(a_i=x)的(i)的个数。
按照暴力的想法我们对(fr_x,bk_x)做个(FFT),然后枚举块中的(j)计算答案。
时间复杂度:(O(frac nS imes nlogn))。
情况(2):三点中存在属于同一块中的点
假设枚举(p,q(p<q))两点属于同一块(另一点(w)可以不在同一块,也可以在同一块,但不能在(p,q)中间)。
我们同样统计(fr_x,bk_x)(我想应该不用再重复一次它们的定义了。。。),于是又分三种情况:
- (w)在之前的块:则(p)对应(j),给答案加上(fr_{2 imes a_p-a_q})。
- (w)在之后的块:则(q)对应(j),给答案加上(bk_{2 imes a_q-a_p})。
- (w)在同一块:显然我们只需考虑(w<p)的情况(若(w>q),则在枚举到(q,w)时会重复计算贡献),然后我们发现实际上只要把(fr_x)改为统计(1sim p-1)的所有数中(a_i=x)的(i)的个数,就可以把这种情况和(w)在之前的块的情况一并计算了,而这个统计显然是可以在枚举(p)的过程中同时维护的。
时间复杂度:(O(nS))。
关于块的大小(S)
我们发现这道题中两种情况下复杂度是不等的。
我们令(frac nS imes 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;
}