【CodeChef COUNTARI】Arithmetic Progressions(FFT)

最新推荐文章于 2019-10-04 03:05:03 发布

AC_alvin

最新推荐文章于 2019-10-04 03:05:03 发布

阅读量442

点赞数

分类专栏： FFT

本文链接：https://blog.csdn.net/u011332631/article/details/46696053

版权

FFT 专栏收录该内容

5 篇文章 0 订阅

订阅专栏

题意：找到一个序列中满足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[N],lcnt[N],icnt[N],tn,a[N];
int blo,tot,Rev[N];
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);
}