BZOJ 3509: [CodeChef] COUNTARI(fft+分块)

最新推荐文章于 2019-02-06 10:33:32 发布

转载最新推荐文章于 2019-02-06 10:33:32 发布 · 118 阅读

0 ·

CC 4.0 BY-SA版权

原文链接：http://www.cnblogs.com/sdfzsyq/p/10023174.html

本文解析了一道算法竞赛题目，探讨如何在给定数组中找到满足特定条件的三元组数量。通过分块和快速傅里叶变换（FFT）优化算法，实现高效解题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

题面

Description

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

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

一行一个整数。
Sample Input
10

3 5 3 6 3 4 10 4 5 2

Sample Output
9

解题思路

　　比较有意思的一道题。首先直接的思路就是对于一个点来说，把它两边卷起来算贡献，但这样复杂度是\(O(n^2logn)\)的，无法接受。考虑分块，我们其实可以把原序列分块，然后每次\(fft\)块两边，这样算出来的贡献块内都可以用。对于块内的贡献可以直接暴力预处理出，预处理就是用两个数组，一个表示左边一个表示右边，然后先枚举块，再枚举两重块内元素。设块的大小为\(T\)，那么就有\(n/T\)块，时间复杂度为\(O(n*T+\dfrac{n}{T}*\sqrt{n*logn})\)，然后两边平衡一下发现\(T\)取\(\sqrt{nlogn}\)时最优。

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<cmath>
#include<algorithm>

using namespace std;
const int MAXN = 100005;
const double Pi = acos(-1);
typedef long long LL;

inline int rd(){
    int x=0,f=1;char ch=getchar();
    while(!isdigit(ch)) {f=ch=='-'?0:1;ch=getchar();}
    while(isdigit(ch))  {x=(x<<1)+(x<<3)+ch-'0';ch=getchar();}
    return f?x:-x;
}

int n,a[MAXN],bl[MAXN],l[MAXN],r[MAXN];
int siz,num,rev[MAXN<<2],limit=1,mx;
int L[30005],R[30005];
LL ans;
struct Complex{
    double x,y;
    Complex(double xx=0,double yy=0){
        x=xx;y=yy;
    }
}f[30005<<2],g[30005<<2];

Complex operator+(const Complex A,const Complex B){return Complex(A.x+B.x,A.y+B.y);}
Complex operator-(const Complex A,const Complex B){return Complex(A.x-B.x,A.y-B.y);}
Complex operator*(const Complex A,const Complex B){return Complex(A.x*B.x-A.y*B.y,A.x*B.y+A.y*B.x);}

inline void fft(Complex *f,int type){
    for(int i=0;i<limit;i++)
        if(i<rev[i]) swap(f[i],f[rev[i]]);
    Complex Wn,w,tmp;int len;
    for(int p=2;p<=limit;p<<=1){
        len=p>>1;Wn=Complex(cos(Pi/len),type*sin(Pi/len));
        for(int k=0;k<limit;k+=p){
            w=Complex(1,0);
            for(int t=k;t<k+len;t++){
                tmp=f[t+len]*w;f[t+len]=f[t]-tmp;
                f[t]=f[t]+tmp;w=w*Wn;
            }
        }
    }
    if(type==-1) for(int i=0;i<limit;i++) f[i].x/=limit;
}

inline int max(int x,int y){
    return x>y?x:y;
}

signed main(){
    n=rd();siz=sqrt(n*log2(n))+1;num=n/siz;if(n%siz) siz++; 
    for(int i=1;i<=n;i++) {a[i]=rd();bl[i]=(i-1)/siz+1;R[a[i]]++;}
    for(int i=1;i<=num;i++) l[i]=(i-1)*siz+1,r[i]=i*siz;r[num]=n;int tmp;
    for(int i=1;i<=num;i++){
        for(int j=l[i];j<=r[i];j++) R[a[j]]--;
        for(int j=l[i];j<=r[i];j++){
            for(int k=j+1;k<=r[i];k++){
                tmp=a[j]*2-a[k];
                if(tmp>=0 && tmp<=30000) ans+=L[tmp];
                tmp=a[k]*2-a[j];
                if(tmp>=0 && tmp<=30000) ans+=R[tmp];
            }
            L[a[j]]++;
        }
    }memset(L,0,sizeof(L));
    for(int i=2;i<num;i++){
        mx=-1;limit=1;
        for(int j=1;j<l[i];j++) mx=max(mx,a[j]),L[a[j]]++;
        for(int j=r[i]+1;j<=n;j++) mx=max(mx,a[j]),R[a[j]]++;
        for(int j=0;j<=mx;j++) f[j].x=L[j],g[j].x=R[j];
        while(limit<=mx*2) limit<<=1;
        for(int j=0;j<limit;j++) rev[j]=(rev[j>>1]>>1)|((j&1)?limit>>1:0);
        fft(f,1);fft(g,1);for(int j=0;j<limit;j++) f[j]=f[j]*g[j];fft(f,-1);
        for(int j=l[i];j<=r[i];j++) ans+=(LL)(f[a[j]<<1].x+0.5);
        for(int j=0;j<=mx;j++) L[j]=R[j]=0;
        for(int j=0;j<limit;j++) f[j].x=f[j].y=g[j].x=g[j].y=0;
    }
    printf("%lld\n",ans);
    return 0;
}

转载于:https://www.cnblogs.com/sdfzsyq/p/10023174.html