CodeChef COUNTARI Arithmetic Progressions（分块 + FFT）

题目

Source

http://vjudge.net/problem/142058

Description

Given N integers A1, A2, …. AN, Dexter wants to know how many ways he can choose three numbers such that they are three consecutive terms of an arithmetic progression.

Meaning that, how many triplets (i, j, k) are there such that 1 ≤ i < j < k ≤ N and Aj - Ai = Ak - Aj.

So the triplets (2, 5, 8), (10, 8, 6), (3, 3, 3) are valid as they are three consecutive terms of an arithmetic
progression. But the triplets (2, 5, 7), (10, 6, 8) are not.

Input

First line of the input contains an integer N (3 ≤ N ≤ 100000). Then the following line contains N space separated integers A1, A2, …, AN and they have values between 1 and 30000 (inclusive).

Output

Output the number of ways to choose a triplet such that they are three consecutive terms of an arithmetic progression.

Sample Input

10
3 5 3 6 3 4 10 4 5 2

Sample Output

9

分析

题目大概说给一个长N的序列A，问有多少个三元组<i,j,k>满足i<j<k且Ai+Ak=2Aj。

i<j<k这个关系不好搞，正解好像是分块：

• 把序列分解成若干块，每一块长度大约为B。分三种情况考虑：

1. 对于三个都在同一块的：枚举各个块，然后通过枚举i和k并更新记录j的信息求出对数。时间复杂度$O(N/B\times B\times B)=O(NB)$。
2. 对于只有两个在同一块的：枚举各个块，并更新记录前面所有块和后面所有块的信息，然后枚举块内的两个数，另一个数可能在块前也可能在块后，这样求出对数。时间复杂度$O(N/B\times B\times B)=O(NB)$。
3. 对于三个数都在不同块的：枚举各个块，并更新记录前面所有块和后面所有块的信息，然后构造多项式用FFT求出前后两边组合成各个和的方案数，通过枚举块内的j即可求出对数。时间复杂度$O(N/B\times 65535\times 16+N/B\times B)$

• 然后就是B大小的设定，我设定的是$5\sqrt N$。
• 最后我连枚举都不会枚举。。还有有个地方只考虑是否大于0，没考虑是否小于等于30000，数组越界，WA了好久。。

代码

#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;
#define INF (1<<30)
#define MAXN 66666
const double PI=acos(-1.0);

struct Complex{
	double real,imag;
	Complex(double _real,double _imag):real(_real),imag(_imag){}
	Complex(){}
	Complex operator+(const Complex &cp) const{
		return Complex(real+cp.real,imag+cp.imag);
	}
	Complex operator-(const Complex &cp) const{
		return Complex(real-cp.real,imag-cp.imag);
	}
	Complex operator*(const Complex &cp) const{
		return Complex(real*cp.real-imag*cp.imag,real*cp.imag+cp.real*imag);
	}
	void setValue(double _real=0,double _imag=0){
		real=_real; imag=_imag;
	}
};

int len;
Complex wn[MAXN],wn_anti[MAXN];

void FFT(Complex y[],int op){
	for(int i=1,j=len>>1,k; i<len-1; ++i){
		if(i<j) swap(y[i],y[j]);
		k=len>>1;
		while(j>=k){
			j-=k;
			k>>=1;
		}
		if(j<k) j+=k;
	}
	for(int h=2; h<=len; h<<=1){
		Complex Wn=(op==1?wn[h]:wn_anti[h]);
		for(int i=0; i<len; i+=h){
			Complex W(1,0);
			for(int j=i; j<i+(h>>1); ++j){
				Complex u=y[j],t=W*y[j+(h>>1)];
				y[j]=u+t;
				y[j+(h>>1)]=u-t;
				W=W*Wn;
			}
		}
	}
	if(op==-1){
		for(int i=0; i<len; ++i) y[i].real/=len;
	}
}
void Convolution(Complex A[],Complex B[],int n){
	for(len=1; len<(n<<1); len<<=1);
	for(int i=n; i<len; ++i){
		A[i].setValue();
		B[i].setValue();
	}

	FFT(A,1); FFT(B,1);
	for(int i=0; i<len; ++i){
		A[i]=A[i]*B[i];
	}
	FFT(A,-1);
}

int a[111111];
int cnt[33100],cnt0[33100],cnt1[33100];
Complex A[MAXN],B[MAXN];

int main(){
	for(int i=0; i<MAXN; ++i){
		wn[i].setValue(cos(2.0*PI/i),sin(2.0*PI/i));
		wn_anti[i].setValue(wn[i].real,-wn[i].imag);
	}
	int n;
	while(~scanf("%d",&n)){
		for(int i=0; i<n; ++i){
			scanf("%d",a+i);
		}
		int block=(int)(sqrt(n)+1e-6)*5;

		long long ans=0;

		for(int b=0; b<n; b+=block){
			for(int i=0; i<block && b+i<n; ++i){
				for(int j=i+1; j<block && b+j<n; ++j){
					int tmp=a[b+i]+a[b+j];
					if((tmp&1)==0){
						ans+=cnt[tmp>>1];
					}
					++cnt[a[b+j]];
				}
				for(int j=i+1; j<block && b+j<n; ++j){
					--cnt[a[b+j]];
				}
			}
		}

		memset(cnt0,0,sizeof(cnt0));
		memset(cnt1,0,sizeof(cnt1));
		for(int i=0; i<n; ++i){
			++cnt1[a[i]];
		}
		for(int b=0; b<n; b+=block){
			for(int i=0; i<block && b+i<n; ++i){
				--cnt1[a[b+i]];
			}
			for(int i=0; i<block && b+i<n; ++i){
				for(int j=i+1; j<block && b+j<n; ++j){
					int tmp=a[b+i]*2-a[b+j];
					if(tmp>0 && tmp<=30000) ans+=cnt0[tmp];
					tmp=a[b+j]*2-a[b+i];
					if(tmp>0 && tmp<=30000) ans+=cnt1[tmp];
				}
			}
			for(int i=0; i<block && b+i<n; ++i){
				++cnt0[a[b+i]];
			}
		}

		memset(cnt0,0,sizeof(cnt0));
		memset(cnt1,0,sizeof(cnt1));
		for(int i=0; i<n; ++i){
			++cnt1[a[i]];
		}
		for(int b=0; b<n; b+=block){
			for(int i=0; i<block && b+i<n; ++i){
				--cnt1[a[b+i]];
			}

			for(int i=0; i<=30000; ++i){
				A[i].setValue(cnt0[i]);
				B[i].setValue(cnt1[i]);
			}
			Convolution(A,B,30001);
			for(int i=0; i<block && b+i<n; ++i){
				ans+=(long long)(A[a[b+i]<<1].real+0.5);
			}

			for(int i=0; i<block && b+i<n; ++i){
				++cnt0[a[b+i]];
			}
		}

		printf("%lld\n",ans);
	}
	return 0;
}