BJ模拟 等差数列（分块+FFT）

Description

      给定N个整数  A1,A2,⋯AN  ，求有多少个三元组  (i,j,k)  满足  1≤i<j<k≤N  且  Aj−Ai=Ak−Aj  。

Input

      第一行一个正整数  N(3≤N≤105)

      第二行  N  个整数  A1,A2,⋯,AN  ，  1≤Ai≤3×104

Output

      输出一个整数表示答案。

Sample Input

10
3 5 3 6 3 4 10 4 5 2

Sample Output

9

HINT

【样例解释】

1 : (i, j, k) = (1, 3, 5), (Ai, Aj, Ak) = (3, 3, 3)
2 : (i, j, k) = (1, 6, 9), (Ai, Aj, Ak) = (3, 4, 5)
3 : (i, j, k) = (1, 8, 9), (Ai, Aj, Ak) = (3, 4, 5)
4 : (i, j, k) = (3, 6, 9), (Ai, Aj, Ak) = (3, 4, 5)
5 : (i, j, k) = (3, 8, 9), (Ai, Aj, Ak) = (3, 4, 5)
6 : (i, j, k) = (4, 6, 10), (Ai, Aj, Ak) = (6, 4, 2)
7 : (i, j, k) = (4, 8, 10), (Ai, Aj, Ak) = (6, 4, 2)
8 : (i, j, k) = (5, 6, 9), (Ai, Aj, Ak) = (3, 4, 5)
9 : (i, j, k) = (5, 8, 9), (Ai, Aj, Ak) = (3, 4, 5)

【数据范围与约定】

对于20%的数据，  N≤5000

对于另外30%的数据，  N≤105,1≤Ai≤50

对于所有数据，  N≤105,1≤Ai≤3×104

题解：分块+FFT

对于30分的部分分：

考虑枚举j。我们可以用pre[t]表示在前j个数中，t这个数出现了多少次，nxt[t]表示后n-j个数中，t这个数出现了多少次。

那么答案就是 ∑u+v=2Ajpre[u]×nxt[v] 。

满分做法：

上面的式子很熟悉吧？就是卷积的形式。不妨将序列分块，pre[t]维护当前块前面，t这个数出现多少次，nxt[t]维护当前块后面，t这个数出现多少次。

那么我们只需要对于每个块的pre和nxt做一次fft，即可快速求出i端点在块前面，j端点在块里面，k端点在块后面的答案。

接下来只需要暴力求一下i, k在块里面的答案即可。

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<string>
#include<algorithm>
#include<cmath>
using namespace std;
typedef unsigned long long ll;
const int Maxn=3e4+50;
const int MaxN=1e5+50;
const double PI=2*acos(-1.0);
inline int read()
{
	char ch=getchar();int i=0,f=1;
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){i=(i<<3)+(i<<1)+ch-'0';ch=getchar();}
	return i*f;
}
struct Complex
{
	double r,i;
	Complex(double r=0,double i=0):r(r),i(i){}
	friend inline Complex operator +(const Complex &a,const Complex &b){return Complex(a.r+b.r,a.i+b.i);}
	friend inline Complex operator -(const Complex &a,const Complex &b){return Complex(a.r-b.r,a.i-b.i);}
	friend inline Complex operator *(const Complex &a,const Complex &b){return Complex(a.r*b.r-a.i*b.i,a.r*b.i+a.i*b.r);}
	inline Complex conj()
	{
		return Complex(r,-i);
	}
}A[Maxn*10],B[Maxn*10],C[Maxn*10];
struct FFT
{
	int pos[Maxn*10],k;
	inline void Init(int len)
	{
		for(k=1;k<=len;k<<=1);
		for(int i=1;i<k;i++)
		pos[i]=(i&1)?((pos[i>>1]>>1)^(k>>1)):(pos[i>>1]>>1);
	}
	inline void fft(Complex *a,int k,int dft)
	{
		for(int i=1;i<k;i++)
		if(i<pos[i])swap(a[i],a[pos[i]]);
		for(int m1=1,m2;m1<k;m1<<=1)
		{
			m2=m1<<1;
			double temp=PI/m2*(double)dft;
			Complex wn(cos(temp),sin(temp));
			for(int i=0;i<k;i+=m2)
			{
				Complex w(1,0);
				for(int j=0;j<m1;j++)
				{
					Complex &A1=a[i+j],&B1=a[i+j+m1],t=w*B1;
					B1=(A1-t);
					A1=(A1+t);
					w=w*wn;
				}
			}
		}
	}
	inline void muitiply(Complex *a,Complex *b,Complex *c)
	{
		fft(a,k,1),fft(b,k,1);
		for(int i=0;i<k;i++)c[i]=a[i]*b[i];
		fft(c,k,-1);
		for(int i=0;i<k;i++)c[i].r/=k;
	}
}fft;
int a[MaxN],r[Maxn*10],l[Maxn*10],sz,m,n,st[MaxN/2000],ed[MaxN/2000],maxa;
ll ans;
int main()
{
	n=read();
	for(int i=1;i<=n;i++)
	{
		a[i]=read();
		r[a[i]]++;
		maxa=max(maxa,a[i]);
	}
	fft.Init(maxa*2+1);
	sz=3500;
	m=(n-1)/sz+1;
	for(int i=1;i<=m;i++)
	{
		st[i]=(i-1)*sz+1;
		ed[i]=i*sz;
	}
	ed[m]=n;
	for(int i=1;i<=m;i++)
	{
		for(int j=st[i];j<=ed[i];j++)r[a[j]]--;
		for(int j=0;j<fft.k;j++)
			A[j]=Complex(l[j],0);
		for(int j=0;j<fft.k;j++)
			B[j]=Complex(r[j],0);
		fft.muitiply(A,B,C);
		for(int j=st[i];j<=ed[i];j++)
		{
			ans+=(ll)(C[2*a[j]].r+0.5);
		}
		for(int p=st[i];p<=ed[i];p++)
		{
			for(int k=st[i];k<p;k++)
			{
				if(2*a[p]-a[k]>0)
				{
					ans+=r[2*a[p]-a[k]];
				}
			}
			for(int k=p+1;k<=ed[i];k++)
			{
				if(2*a[p]-a[k]>0)
				{
					ans+=l[2*a[p]-a[k]];
				}
			}
			l[a[p]]++;
		}
	}
	cout<<ans<<endl;
}