BZOJ-3509 母函数+分块+暴力+FFT

题目描述

给定一个长度为N的数组A[]，求有多少对i, j, k（1<=i<j<k<=N）满足A[k]-A[j]=A[j]-A[i]。

---

输入格式

第一行一个整数N（N<=10^5）。
接下来一行N个数A[i]（A[i]<=30000）。

---

输出格式

一行一个整数。

---

样例输入

10
3 5 3 6 3 4 10 4 5 2

---

样例输出

9

这题网上题解很多
我的第一个想法是枚举 j 然后左边的数 和 右边的数 FFT处理一下 
这样的复杂度是  n^2*log（n） 复杂度明显不行 所以需要进一步的优化
网上题解都是用分块进行优化
先在每一个块里面找一下是否存在满足条件的
然后再找块外面的满足条件的
这题BZOJ 开了40S 跑的了的
我FFT是直接套板子  当作黑盒使用 所有代码有点鬼畜

#include <cstdio>
#include <cstring>
#include <queue>
#include <cmath>
#include <algorithm>
#include <set>
#include <iostream>
#include <map>
#include <stack>
#include <string>
#include <vector>
#define  pi acos(-1.0)
#define  eps 1e-9
#define  fi first
#define  se second
#define  rtl   rt<<1
#define  rtr   rt<<1|1
#define  bug         printf("******\n")
#define  mem(a,b)    memset(a,b,sizeof(a))
#define  name2str(x) #x
#define  fuck(x)     cout<<#x" = "<<x<<endl
#define  f(a)        a*a
#define  sf(n)       scanf("%d", &n)
#define  sff(a,b)    scanf("%d %d", &a, &b)
#define  sfff(a,b,c) scanf("%d %d %d", &a, &b, &c)
#define  sffff(a,b,c,d) scanf("%d %d %d %d", &a, &b, &c, &d)
#define  pf          printf
#define  FRE(i,a,b)  for(i = a; i <= b; i++)
#define  FREE(i,a,b) for(i = a; i >= b; i--)
#define  FRL(i,a,b)  for(i = a; i < b; i++)+
#define  FRLL(i,a,b) for(i = a; i > b; i--)
#define  FIN         freopen("data.txt","r",stdin)
#define  gcd(a,b)    __gcd(a,b)
#define  lowbit(x)   x&-x
#define rep(i,a,b) for(int i=a;i<b;++i)
#define per(i,a,b) for(int i=a-1;i>=b;--i)
using namespace std;
typedef long long  LL;
typedef unsigned long long ULL;
const int maxn = 3e5 + ;
const int maxm = maxn * ;
int n, m, a[maxn], b[maxn];
int len, res[maxm], mx; //开大4倍
struct cpx {
    long double r, i;
    cpx ( long double r = , long double i =  ) : r ( r ), i ( i ) {};
    cpx operator+ ( const cpx &b ) {
        return cpx ( r + b.r, i + b.i );
    }
    cpx operator- ( const cpx &b ) {
        return cpx ( r - b.r, i - b.i );
    }
    cpx operator* ( const cpx &b ) {
        return cpx ( r * b.r - i * b.i, i * b.r + r * b.i );
    }
} va[maxm], vb[maxm];
void rader ( cpx F[], int len ) { //len = 2^M,reverse F[i] with  F[j] j为i二进制反转
    int j = len >> ;
    for ( int i = ; i < len - ; ++i ) {
        if ( i < j ) swap ( F[i], F[j] ); // reverse
        int k = len >> ;
        while ( j >= k ) j -= k, k >>= ;
        if ( j < k ) j += k;
    }
}
void FFT ( cpx F[], int len, int t ) {
    rader ( F, len );
    for ( int h = ; h <= len; h <<=  ) {
        cpx wn ( cos ( -t *  * pi / h ), sin ( -t *  * pi / h ) );
        for ( int j = ; j < len; j += h ) {
            cpx E ( ,  ); //旋转因子
            for ( int k = j; k < j + h / ; ++k ) {
                cpx u = F[k];
                cpx v = E * F[k + h / ];
                F[k] = u + v;
                F[k + h / ] = u - v;
                E = E * wn;
            }
        }
    }
    if ( t == - ) //IDFT
        for ( int i = ; i < len; ++i ) F[i].r /= len;
}
void Conv ( cpx a[], cpx b[], int len ) { //求卷积
    FFT ( a, len,  );
    FFT ( b, len,  );
    for ( int i = ; i < len; ++i ) a[i] = a[i] * b[i];
    FFT ( a, len, - );
}
void gao () {
    len = ;
    mx = n + m;
    while ( len <= mx ) len <<= ; //mx为卷积后最大下标
    for ( int i = ; i < len; i++ ) va[i].r = va[i].i = vb[i].r = vb[i].i = ;
    for ( int i = ; i < n; i++ ) va[i].r = a[i]; //根据题目要求改写
    for ( int i = ; i < m; i++ ) vb[i].r = b[i]; //根据题目要求改写
    Conv ( va, vb, len );
    for ( int i = ; i < len; ++i ) res[i] += ( LL ) floor ( va[i].r + 0.5 );
}
int L[maxn], R[maxn], p[maxn];
int main() {
    //FIN;
    int num;
    sf ( num );
    m = -;
    for ( int i = ; i < num ; i++ ) {
        sf ( p[i] );
        R[p[i]]++;
        m = max ( m, p[i] );
    }
    m++;
    n = m;
    int sz = min ( num,  * ( int ) sqrt ( num ) );
    LL ans = ;
    for ( int i =  ; i < num ; i += sz ) {
        int r = min ( num, i + sz ) - ;
        for ( int j = i ; j <= r ; j++ ) R[p[j]]--;
        for ( int j = i, x ; j <= r ; j++ ) {
            for ( int k = j +  ; k <= r ; k++ ) {
                x =  * p[j] - p[k];
                if ( x >=  ) ans += 1LL * L[x];
                x =  * p[k] - p[j];
                if ( x >=  ) ans += 1LL * R[x];
            }
            L[p[j]]++;
        }
    }
    for ( int i =  ; i < num ; i += sz ) {
        int r = min ( num, i + sz ) - ;
        mem ( a,  ), mem ( b,  ), mem ( res,  );
        for ( int j =  ; j < i ; j++ ) a[p[j]]++;
        for ( int j = r +  ; j < num ; j++ ) b[p[j]]++;
        gao();
        for ( int j = i ; j <= r ; j++ ) ans += 1LL * ( res[ * p[j]] );
    }
    printf ( "%lld\n", ans );
    return ;
}