【CodeChef COUNTARI】Arithmetic Progressions(FFT)

题意：找到一个序列中满足i < j < k && Ai + Ak = 2 * Aj的方案数

这道题可以直接用分块来做，假设每一块的大小是C，现在考虑第j块，那么可以分为三种情况

1、对于三个数都在块j中，则可以通过枚举两个数找第三个数O（C^2）

2、对于两个数在块j中，也可以通过枚举两个数找第三个数O(C^2)

3、对于只有一个在块j中，则是通过对j - 1个块和后面n - j个块做FFT，就可以得到答案。

具体见代码：

#include<cstdio>
#include<cmath>
#include<vector>
#include<cstring>
#include<algorithm>
using namespace std;
const int N = 200005;
const double pi = acos(-1.0);
typedef long long LL;
struct Complex{
double r,i;
Complex(){}
Complex(double _r,double _i):r(_r),i(_i){}
Complex operator + (const Complex &a)const{
return Complex(r + a.r,i + a.i);
}
Complex operator - (const Complex &a)const{
return Complex(r - a.r,i - a.i);
}
Complex operator * (const Complex &a)const{
return Complex(r * a.r - i * a.i,r * a.i + i * a.r);
}
Complex operator / (const double &n)const{
return Complex(r / n,i / n);
}
};
Complex x1[N << 1],x2[N << 1],tmp[N << 1];
int scnt
,lcnt
,icnt
,tn,a
;
int blo,tot,Rev
;
int rev(int x){
int res = 0;
for(int i = 0 ; i < tn ; i ++){
if(x & 1) res += (1 << tn - 1 - i);
x >>= 1;
}
return res;
}
void fft(Complex A[],int n,int op){
for(int i = 0 ; i < n ; i ++) tmp[ Rev[i] ] = A[i];
for(int i = 0 ; i < n ; i ++) A[i] = tmp[i];
for(int i = 1 ; (1 << i) <= n ; i ++){
int m = 1 << i;
Complex wn(cos(op * 2 * pi / m),sin(op * 2 * pi / m));
for(int k = 0 ; k < n ; k += m){
Complex w(1,0),u,t;
for(int j = 0 ; j < m / 2 ; j ++){
u = A[k + j];
t = w * A[k + j + m / 2];
A[k + j] = u + t;
A[k + j + m / 2] = u - t;
w = w * wn;
}
}
}
if(op == -1) for(int i = 0 ; i < n ; i ++)
A[i] = A[i] / n;
}
int init(int n,int Max){
int len;
blo = min(35,n);
tot = (n - 1 + blo) / blo;
tn = ceil(log(Max + 1.0) / log(2.0)) + 1;
len = 1 << tn;
for(int i = 0 ; i < len ; i ++) Rev[i] = rev(i);
return len;
}
void solve(int n){
int Max = 0,len;
for(int i = 0 ; i < n ; i ++){
scanf("%d",&a[i]);
Max = max(Max,a[i]);
lcnt[ a[i] ] ++;
}
len = init(n,Max);
LL res = 0;
for(int t = 0 ; t < blo ; t ++){
int l,r;
l = t * tot;r = min((t + 1) * tot,n);
for(int i = l ; i < r ; i ++) lcnt[ a[i] ] --;
for(int i = 0 ; i < len ; i ++){
x1[i] = Complex(scnt[i],0);
x2[i] = Complex(lcnt[i],0);
}
fft(x1,len,1);
fft(x2,len,1);
for(int i = 0 ; i < len ; i ++)
x1[i] = x1[i] * x2[i];
fft(x1,len,-1);
for(int i = l ; i < r ; i ++)
res += LL(x1[ a[i] * 2 ].r + 0.5);
//*****************************************************
for(int i = l ; i < r ; i ++){
for(int j = i + 1 ; j < r ; j ++){
int tt = a[j] * 2 - a[i];
if(tt >= 1 && tt < N) res += lcnt[tt];
tt = a[i] * 2 - a[j];
if(tt >= 1 && tt < N) res += icnt[tt] + scnt[tt];
}
icnt[ a[i] ] ++;
}
for(int i = l ; i < r ; i ++){
scnt[ a[i] ] ++;
icnt[ a[i] ] --;
}
}
printf("%lld\n",res);
}
int main()
{
int n;
while(~scanf("%d",&n)) solve(n);
}