【分块】【FFT】CodeChef COUNTARI Arithmetic Progressions

分析:

主要就是分块。。。

首先,假设我们将原序列分为k块,

对于任意一个三元组,如果三个都在同一个块内,或者两个在一个块内,都可以在 O ( N k × N k × k ) = O ( N 2 K ) 的复杂度内解决。

现在考虑只有三个值在不同的块的情况。

这样对于同一个块,可以把之前的块的每种数的个数存入一个多项式,之后的块中每个数的个数存入一个多项式。然后现在无非就是枚举这个块内的值 a i ,求两个多项式卷积后值为 2 a i 的个数。(同一个块只FFT一次)

这部分的复杂度就是 O ( k × 65536 × l o g 2 ( 65536 ) + N )
总的复杂度就是 O ( N 2 k + k 1048576 )

然后根据均值不等式,求得当N=k*1024时取最小,最小值为 N × 2048 勉强还是卡得过。

所以块的大小就是1024(这里不直接令k的大小,而是块的大小,因为块的大小为定值,当然其实直接令k=100也没什么区别)

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<vector>
#include<set>
#define SF scanf
#define PF printf
#define MAXN 100000
using namespace std;
typedef long long ll;
int siz=65536;
const double Pi=acos(-1);
struct cpx{
    double r,i;
    cpx() {}
    cpx(double _r,double _i):r(_r),i(_i) {}
    cpx operator * (const cpx &a) const{
        return cpx(r*a.r-i*a.i,r*a.i+i*a.r);
    }
    cpx operator + (const cpx &a) const{
        return cpx(r+a.r,i+a.i);
    }
    cpx operator - (const cpx &a) const{
        return cpx(r-a.r,i-a.i);
    }
};
void fft(cpx *a,int f,int N){
    int i,j,k;
    for(i=1,j=0;i<N;i++){
        for(int d=N;j^=d>>=1,~j&d;);
        if(i<j)
            swap(a[i],a[j]);
    }
    for(i=1;i<N;i<<=1){
        cpx wn(cos(Pi/i),f*sin(Pi/i));
        for(j=0;j<N;j+=i<<1){
            cpx w(1,0);
            for(k=0;k<i;k++,w=w*wn){
                cpx x=a[j+k],y=w*a[i+j+k];
                a[j+k]=x+y;
                a[i+j+k]=x-y;
            }
        }
    }
    if(f==-1)
        for(i=0;i<N;i++)
            a[i].r/=N;
}
cpx A[MAXN],B[MAXN];
int sum[MAXN],n,p[MAXN];
ll ans;
int main(){
    SF("%d",&n);
    int k=1024;
    for(int i=0;i<n;i++)
        SF("%d",&p[i]);
    for(int i=0;i<n;i+=k){
        memset(sum,0,sizeof sum);
        for(int j=0;j<k&&i+j<n;j++){
            for(int l=j+1;l<k&&i+l<n;l++){
                int x=p[i+j]*2-p[i+l];
                if(x>=0&&x<=30000)
                    ans+=sum[x];    
            }
            sum[p[i+j]]++;
        }
    }
    memset(sum,0,sizeof sum);
    for(int i=0;i<n;i+=k){
        for(int j=0;j<k&&i+j<n;j++)
            for(int l=j+1;l<k&&i+l<n;l++){
                int x=p[i+j]*2-p[i+l];
                if(x>=0&&x<=30000)
                    ans+=sum[x];    
            }
        for(int j=0;j<k&&i+j<n;j++)
            sum[p[i+j]]++;
    }
    memset(sum,0,sizeof sum);
    for(int i=((n-1)/k)*k;i>=0;i-=k){
        for(int j=0;j<k&&i+j<n;j++)
            for(int l=j-1;l>=0;l--){
                int x=p[i+j]*2-p[i+l];
                if(x>=0&&x<=30000)
                    ans+=sum[x];    
            }
        for(int j=0;j<k&&i+j<n;j++)
            sum[p[i+j]]++;
    }
    for(int i=0;i<n;i+=k){
        memset(A,0,sizeof A);
        memset(B,0,sizeof B);
        for(int j=0;j<i;j++)
            A[p[j]].r++;
        for(int j=i+k;j<n;j++)
            B[p[j]].r++;
        fft(A,1,siz);
        fft(B,1,siz);
        for(int j=0;j<siz;j++)
            A[j]=A[j]*B[j];
        fft(A,-1,siz);
        for(int j=i;j<i+k&&j<n;j++)
            ans+=(long long)(A[p[j]*2].r+0.5);
    }
    PF("%lld",ans);
}

猜你喜欢

转载自blog.csdn.net/qq_34454069/article/details/80795662