Arithmetic Progressions CC COUNTARI

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <queue>
#include <algorithm>
#include <vector>
#include <cstring>
#include <stack>
#include <cctype>
#include <utility>   
#include <map>
#include <string>  
#include <climits> 
#include <set>
#include <string>    
#include <sstream>
#include <utility>   
#include <ctime>
#include <bitset>
//#pragma comment(linker, "/STACK:102400000,102400000")

using std::priority_queue;
using std::vector;
using std::swap;
using std::stack;
using std::sort;
using std::max;
using std::min;
using std::pair;
using std::map;
using std::string;
using std::cin;
using std::cout;
using std::set;
using std::queue;
using std::string;
using std::stringstream;
using std::make_pair;
using std::getline;
using std::greater;
using std::endl;
using std::multimap;
using std::deque;
using std::unique;
using std::lower_bound;
using std::random_shuffle;
using std::bitset;
using std::upper_bound;
using std::multiset;

typedef long long LL;
typedef unsigned long long ULL;
typedef unsigned UN;
typedef pair<int, int> PAIR;
typedef multimap<int, int> MMAP;
typedef LL TY;
typedef long double LF;

const int MAXN(1500010);
const int MAXM(50010);
const int MAXE(150010);
const int MAXK(6);
const int HSIZE(13131);
const int SIGMA_SIZE(4);
const int MAXH(20);
const int INFI((INT_MAX-1) >> 1);
const ULL BASE(31);
const LL LIM(1e13);
const int INV(-10000);
const int MOD(31313);
const double EPS(1e-7);
const LF PI(acos(-1.0));

template<typename T> inline void checkmax(T &a, T b){if(b > a) a = b;}
template<typename T> inline void checkmin(T &a, T b){if(b < a) a = b;}
template<typename T> inline T ABS(const T &a){return a < 0? -a: a;}

struct COMPLEX
{
	double a, b;
	COMPLEX(double a_, double b_): a(a_), b(b_){}
	COMPLEX(){}
	friend COMPLEX operator +(const COMPLEX &op1, const COMPLEX &op2){ return COMPLEX(op1.a+op2.a, op1.b+op2.b);}
	friend COMPLEX operator -(const COMPLEX &op1, const COMPLEX &op2){ return COMPLEX(op1.a-op2.a, op1.b-op2.b);}
	friend COMPLEX operator *(const COMPLEX &op1, const COMPLEX &op2){ return COMPLEX(op1.a*op2.a-op1.b*op2.b, op1.a*op2.b+op1.b*op2.a);}
};

void change(COMPLEX *y, int len)
{
	for(int i = 1, j = len/2; i < len-1; ++i)
	{
		if(i < j) swap(y[i], y[j]);
		int k = len/2;
		while(j >= k)
		{
			j -= k;
			k /= 2;
		}
		if(j < k) j += k;
	}
}

void FFT(COMPLEX *y, int len, int flag)
{
	change(y, len);
	for(int h = 2; h <= len; h <<= 1)
	{
		COMPLEX wn(cos(flag*2.0*PI/h), sin(flag*2.0*PI/h));
		for(int j = 0; j < len; j += h)
		{
			COMPLEX w(1, 0);
			int temp = h >> 1;
			for(int k = j; k < j+temp; ++k)
			{
				COMPLEX u = y[k];
				COMPLEX t = w*y[k+temp];
				y[k] = u+t;
				y[k+temp] = u-t;
				w = w*wn;
			}
		}
	}
	if(flag == -1)
		for(int i = 0; i < len; ++i) y[i].a /= len;
}

int arr[100010];
int cnt1[60010], cnt2[60010], cnt3[60010];
COMPLEX x1[70010], x2[70010];

int main()
{
	int n;
	while(~scanf("%d", &n))
	{
		int mx = 0;
		for(int i = 0; i < n; ++i)
		{
			scanf("%d", arr+i);
			checkmax(mx, arr[i]);
		}
		int K = sqrt(n/(log((double)n)/log(2.)))+1;
		int tn = max(2, n/K+1);
		LL ans = 0;
		int ind = 0;
		memset(cnt1, 0, sizeof(cnt1));  //区间左边
		memset(cnt2, 0, sizeof(cnt2));  //区间右边
		for(int i = 0; i < n; ++i) ++cnt2[arr[i]];
		for(int i = 0; i < n; i += tn)
		{
			memset(cnt3, 0, sizeof(cnt3));  //区间内
			for(int ti = i, j = 0; j < tn && ti < n; ++ti, ++j) --cnt2[arr[ti]];
			for(int ti = i, j = 0; j < tn && ti < n; ++ti, ++j)
			{
				for(int tti = ti+1, tj = j+1; tj < tn && tti < n; ++tti, ++tj)
				{
					ans += 2*arr[ti]-arr[tti] >= 0? cnt1[2*arr[ti]-arr[tti]]: 0;
					ans += 2*arr[tti]-arr[ti] >= 0? cnt2[2*arr[tti]-arr[ti]]: 0;
					ans += 2*arr[ti]-arr[tti] >= 0? cnt3[2*arr[ti]-arr[tti]]: 0;
				}
				++cnt3[arr[ti]];
			}
			for(int j = 0; j <= mx; ++j)
			{
				x1[j] = COMPLEX(cnt1[j], 0);
				x2[j] = COMPLEX(cnt2[j], 0);
			}
			int temp = (mx+1)*2, l = 1;
			while(l < temp) l <<= 1;
			for(int j = mx+1; j < l; ++j) x1[j] = x2[j] = COMPLEX(0, 0);
			FFT(x1, l, 1);
			FFT(x2, l, 1);
			for(int j = 0; j < l; ++j) x1[j] = x1[j]*x2[j];
			FFT(x1, l, -1);
			for(int ti = i, j = 0; j < tn && ti < n; ++ti, ++j) ans += (LL)(x1[arr[ti]*2].a+0.5);
			for(int ti = i, j = 0; j < tn && ti < n; ++ti, ++j) ++cnt1[arr[ti]];
		}
		printf("%lld\n", ans);
	}
	return 0;
}