bzoj 3509: [CodeChef] COUNTARI 分块fft

最新推荐文章于 2020-08-01 21:08:18 发布

_beginend

最新推荐文章于 2020-08-01 21:08:18 发布

阅读量322

点赞数 1

分类专栏：分块快速傅里叶变换

本文链接：https://blog.csdn.net/qq_33229466/article/details/79401968

版权

快速傅里叶变换同时被 2 个专栏收录

40 篇文章 0 订阅

订阅专栏

分块

36 篇文章 0 订阅

订阅专栏

题意

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

分析

把式子画一下就变成了A[j]*2=A[k]+A[i]，一个很显然的想法就是枚举j然后fft，复杂度是O(NVlogV)，显然不能接受。
考虑分块，设块的大小为B，我们可以枚举每一块，然后分两种情况，一种是i和k都不在块内，这时可以用一次fft来统计完整块的贡献；另一种是i和k至少有一个在块内，我们可以枚举j个其中一个，通过O(B^2)来计算块内的贡献。
这样做的复杂度是O(N/B*VlogV+B^2)，解一下方程可以得到B大概等于3560左右，然后就可以跑过去了。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;

typedef long long LL;

const int B=3560;
const double pi=acos(-1);
const int N=100005;
const int M=30005;

int n,w[N],pre[M*2],suf[M*2],rev[M*4],V,L,bel[N],sta[M],end[M];

struct com
{
    double x,y;
    com operator + (const com &a) const {return (com){x+a.x,y+a.y};}
    com operator - (const com &a) const {return (com){x-a.x,y-a.y};}
    com operator * (const com &a) const {return (com){x*a.x-y*a.y,x*a.y+y*a.x};}
    com operator / (const double &a) const {return (com){x/a,y/a};}
}a[M*4],b[M*4];

int read()
{
    int x=0,f=1;char ch=getchar();
    while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}

void fft(com *a,int f)
{
    for (int i=0;i<L;i++) if (i<rev[i]) swap(a[i],a[rev[i]]);
    for (int i=1;i<L;i<<=1)
    {
        com wn=(com){cos(pi/i),f*sin(pi/i)};
        for (int j=0;j<L;j+=(i<<1))
        {
            com w=(com){1,0};
            for (int k=0;k<i;k++)
            {
                com u=a[j+k],v=a[j+k+i]*w;
                a[j+k]=u+v;a[j+k+i]=u-v;
                w=w*wn;
            }
        }
    }
    if (f==-1) for (int i=0;i<L;i++) a[i]=a[i]/L;
}

int main()
{
    n=read();
    for (int i=1;i<=n;i++)
    {
        w[i]=read();bel[i]=(i+B-1)/B;suf[w[i]]++;V=max(V,w[i]);
        if (!sta[bel[i]]) sta[bel[i]]=i;
        end[bel[i]]=i;
    }
    int lg=0;
    for (L=1;L<=V*2;L<<=1,lg++);
    for (int i=0;i<L;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(lg-1));
    LL ans=0;
    for (int i=1;i<=bel[n];i++)
    {
        int l=sta[i],r=end[i];
        for (int j=l;j<=r;j++) suf[w[j]]--;
        for (int j=0;j<L;j++) a[j]=b[j]=(com){0,0};
        for (int j=0;j<=V;j++) a[j]=(com){pre[j],0},b[j]=(com){suf[j],0};
        fft(a,1);fft(b,1);
        for (int j=0;j<L;j++) a[j]=a[j]*b[j];
        fft(a,-1);
        for (int j=l;j<=r;j++) ans+=(LL)(a[w[j]*2].x+0.1);
        for (int j=l;j<=r;j++)
        {
            for (int k=j+1;k<=r;k++) ans+=(w[j]*2-w[k]>=0)?pre[w[j]*2-w[k]]:0;
            pre[w[j]]++;
        }
        for (int j=l+1;j<=r;j++)
            for (int k=l;k<j;k++)
                if (w[j]*2-w[k]>=0) ans+=suf[w[j]*2-w[k]];
    }
    printf("%lld",ans);
    return 0;
}