BZOJ3509: [CodeChef] COUNTARI

BZOJ3509: [CodeChef] COUNTARI

https://lydsy.com/JudgeOnline/problem.php?id=3509

分析:

• ​分块，考虑\(i,j,k\)在不同的块内的方案数，这个我们枚举中间那个块两边做一次\(fft\)再枚举中间块即可。
• 其他的方案数可以用前缀桶和后缀桶求出。

代码:

 #include <cstdio>
#include <cstring>
#include <algorithm>
#include <cstdlib>
#include <cmath>
using namespace std;
typedef long long ll;
typedef double f2;
#define N 200050
#define M 30000
int n,a[N],L[N],R[N],c[2][N];
const f2 pi=acos(-1);
ll ans;
struct cp {
	f2 x,y; 
	cp() {}
	cp(f2 x_,f2 y_) {x=x_, y=y_;}
	cp operator + (const cp &u) const {
		return cp(x+u.x, y+u.y);
	}
	cp operator - (const cp &u) const {
		return cp(x-u.x, y-u.y);
	}
	cp operator * (const cp &u) const {
		return cp(x*u.x-y*u.y,x*u.y+y*u.x);
	}
}A[N],B[N];
void fft(cp *a,int len,int flg) {
	int i,j,k,t; cp w,wn,tmp;
	for(i=k=0;i<len;i++) {
		if(i>k) swap(a[i],a[k]);
		for(j=len>>1;(k^=j)<j;j>>=1) ;
	}
	for(k=2;k<=len;k<<=1) {
		t=k>>1;
		wn=cp(cos(2*pi*flg/k),sin(2*pi*flg/k));
		for(i=0;i<len;i+=k) {
			w=cp(1,0);
			for(j=i;j<i+t;j++) {
				tmp=a[j+t]*w;
				a[j+t]=a[j]-tmp;
				a[j]=a[j]+tmp;
				w=w*wn;
			}
		}
	}
	if(flg==-1) for(i=0;i<len;i++) a[i].x/=len;
}
int main() {
	int size=2000;
	scanf("%d",&n);
	int i,j,k;
	for(i=1;i<=n;i++) scanf("%d",&a[i]),c[1][a[i]]++;
	int blo=(n+size-1)/size;
	for(i=1;i<=blo;i++) {
		L[i]=R[i-1]+1, R[i]=min(n,i*size);
		for(j=L[i];j<=R[i];j++) {
			c[1][a[j]]--;
			for(k=L[i];k<j;k++) {
				if((a[j]<<1)>=a[k]) {
					ans+=c[1][(a[j]<<1)-a[k]];
				}
				if((a[k]<<1)>=a[j]) {
					ans+=c[0][(a[k]<<1)-a[j]];
				}
			}
		}
		for(j=L[i];j<=R[i];j++) c[0][a[j]]++;
	}
	int len=1;
	while(len<=(M<<1)) len<<=1;
	
	memset(c,0,sizeof(c));
	for(i=1;i<=n;i++) c[1][a[i]]++;
	for(i=L[1];i<=R[1];i++) c[0][a[i]]++,c[1][a[i]]--;
	for(i=2;i<blo;i++) {
		for(j=L[i];j<=R[i];j++) c[1][a[j]]--;
		for(j=0;j<len;j++) A[j]=B[j]=cp(0,0);
		for(j=0;j<=M;j++) A[j].x=c[0][j], B[j].x=c[1][j];
		fft(A,len,1), fft(B,len,1);
		for(j=0;j<len;j++) A[j]=A[j]*B[j];
		fft(A,len,-1);
		for(j=L[i];j<=R[i];j++) {
			ans+=int(A[a[j]<<1].x+0.1);
		}
		for(j=L[i];j<=R[i];j++) c[0][a[j]]++;
	}
	printf("%lld\n",ans);
}