[生成函数 FFT 分块] BZOJ 3509 [CodeChef] COUNTARI

学着大神优化了FFT常数：http://blog.csdn.net/geotcbrl/article/details/50636401

可以看出来这是个生成函数

暴力的话是N次FFT 

然后考虑分块 B次FFT + 块内统计

#include<cstdio>
#include<cstdlib>
#include<complex>
#include<algorithm>
#include<cstring>
#define cl(x) memset(x,0,sizeof(x))
#define PI acos(-1.0)
using namespace std;
typedef long long ll;

inline char nc()
{
	static char buf[100000],*p1=buf,*p2=buf;
	if (p1==p2) { p2=(p1=buf)+fread(buf,1,100000,stdin); if (p1==p2) return EOF; }
	return *p1++;
}

inline void read(int &x){
	char c=nc(),b=1;
	for (;!(c>='0' && c<='9');c=nc()) if (c=='-') b=-1;
	for (x=0;c>='0' && c<='9';x=x*10+c-'0',c=nc()); x*=b;
}

struct comp {
    long double real , imag;
    comp(long double real = 0 , long double imag = 0):real(real) , imag(imag) { }
    inline friend comp operator+(comp&a , comp&b)
        { return comp(a.real + b.real , a.imag + b.imag); }
    inline friend comp operator-(comp&a , comp&b)
        { return comp(a.real - b.real , a.imag - b.imag); }
    inline friend comp operator*(comp&a , comp&b)
        { return comp(a.real * b.real - a.imag * b.imag , a.imag * b.real + a.real * b.imag); }
    inline friend void swap(comp&a , comp&b)
        { comp c = a ; a = b ; b = c; }
};

int n,m,Len;
comp a[524299];
ll c[524299],Rev[524299];

inline void FFT(comp *a,int r)
{
	for (int i=0;i<n;i++) if (i<Rev[i]) swap(a[i],a[Rev[i]]);
	for (int i=1;i<n;i<<=1)
	{
		comp wn(cos(PI/i),r*sin(PI/i));
		for (int j=0;j<n;j+=(i<<1))
		{
			comp w(1,0);
			for (int k=0;k<i;k++,w=w*wn)
			{
				comp x=a[j+k],y=w*a[j+i+k];
				a[j+k]=x+y; a[j+i+k]=x-y;				
			}
		}
	}	
	if (r==-1) for (int i=0;i<n;i++) a[i].real/=n,a[i].imag/=n;;
}

int N,B,tot,V;
int A[100005],pos[100005];
int cnt[85][30005],st[85],ed[85];

ll ans;
int L[30005],R[30005],pre[30005];

inline void FFT_init()
{
	m=2*V;
	for (n=1;n<=m;n<<=1) Len++;
	for(int i=0;i<n;i++) Rev[i]=(Rev[i>>1]>>1)|((i&1)<<(Len-1));
}

inline void Mul()
{
//	cl(a);
//	cl(b);
	memset(a,0,sizeof(comp)*(n+1));
	for (int i=0;i<=V;i++) a[i]=comp(L[i]+R[i],L[i]-R[i]);
//	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<n;i++) a[i]*=b[i];
	for (int i=0;i<n;i++) a[i]=a[i]*a[i];
	FFT(a,-1);
	for(int i=0;i<=m;i+=2) c[i]=(ll)(a[i].real/4+0.1);
}

inline void calc(int x)
{
	int t=A[x]+A[x],p=pos[x];
	for (int i=st[p];i<x;i++)
		if (t-A[i]>=0 && t-A[i]<=V)
			ans+=R[t-A[i]];
	for (int i=x+1;i<=ed[p];i++)
		if (t-A[i]>=0 && t-A[i]<=V)
		{
			ans+=L[t-A[i]];
			ans+=pre[t-A[i]];
		}
	if (p!=1 && p!=tot)	ans+=c[t];
}

inline void Solve()
{
	FFT_init();
	for (int i=1;i<=tot;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!=tot)
			Mul();
		cl(pre);
		for (int j=st[i];j<=ed[i];j++)
			calc(j),pre[A[j]]++;
	}
}

int main()
{
	freopen("t.in","r",stdin);
	freopen("t.out","w",stdout);
	read(N); B=2000;
	for (int i=1;i<=N;i++)
	{
		read(A[i]); R[A[i]]++; V=max(V,A[i]);
		pos[i]=(i-1)/B+1;
	}
	tot=pos[N];
	for (int i=1;i<=tot;i++) st[i]=(i-1)*B+1,ed[i]=i*B;
	ed[tot]=min(ed[tot],N);
	Solve();
	printf("%lld\n",ans);
	return 0;
}