[CodeChef] COUNTARI

問題描述

給定一個長度爲N的數組A[]，求有多少對i, j, k（1<=i

輸入格式

第一行一個整數N（N<=10^5）。
接下來一行N個數A[i]（A[i]<=30000）。

輸出格式

一行一個整數。

樣例輸入

10
3 5 3 6 3 4 10 4 5 2

樣例輸出

9

---

題解

PS：之後講的“中間數”即爲A[k]-A[j]=A[j]-A[i]時的A[j]
顯然，對於每個數A[i]，如果我們把A[1]至A[i-1]的值放入大小爲30000的數組c1[]中（即，c1[j]記錄前i-1個數中有多少個值爲j的數），把A[i+1]至A[n]的值放入c2[]中，那麼，c1和c2進行FFT之後，c3[2*A[i]]即爲以i 爲中間數的答案（因爲c3[2*a[i]]=∑$\sum$ c1[l]*c2[k]（l+k==2*A[i]））。
這個東東很優秀。我們可以把時間複雜度由O（n2${}^2$ ）升至O（n2${}^2$log2n$lo{g}_{2}n$ ）。。。
想想，爲什麼進行優秀的FFT之後會變慢？因爲每次FFT之後我們都只用了一個位置的數！！！
也就是說，我們要在每次FFT之後多獲取一點信息。
在老闆的提示下，使用分塊。
我們定每塊的長度爲s，將數列分爲n/s塊。
對於每個數a[i]，答案可能的區域被分爲：

1表示塊外左邊，2表示塊內左邊，3表示塊內右邊，4表示塊外右邊。
所以，將A[i]設爲中間數，另外兩數的範圍有（1,3）、（1,4）、（2,3）、（2,4）。
（1,4）可以用FFT搞，其它三個都可以用正常分塊思想搞出來。
總時間複雜度O（n*s+Max_vlog2MaxV$lo{g}_{2}Ma{x}_V$ ）。但實際上FFT常數較大，每塊長度最好大一點。
PS：坑？
1、自己造的：FFT一定要好好打；
2、不要開二維數組，一個complex_double的空間是16B。。。會G

---

代碼

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <cmath>
#include <complex>
#define db double
#define cd complex<db>
#define ll long long
using namespace std;
const double pi=3.1415926535897932;
cd w[70001];
ll n,s;
void fly(cd a[],ll f){
    ll i,j,k,m;
    for(i=j=0;i<n;i++){
        if(i<j)swap(a[i],a[j]);
        for(k=n>>1;(j^=k)<k;k>>=1);
    }
    w[0]=1;
    for(m=1;m<n;m<<=1){
        cd ha=exp(cd(0,pi*(db)f/(db)m));
        for(i=1;i<m;i++)w[i]=w[i-1]*ha;
        for(i=0;i<n;i+=(m<<1))
            for(j=0;j<m;j++){
                cd p=a[i+j],q=a[i+j+m]*w[j];
                a[i+j]=p+q,a[i+j+m]=p-q;
            }
    }
    if(f==1)return;
    cd tmp=1.0/(db)n;
    for(i=0;i<n;i++)
        a[i]*=tmp;
}
ll gg(cd x)
{return (ll)floor(x.real()+0.5);}
cd ti[70001];
void FFT(cd a[],cd b[],cd c[]){
    for(ll i=0;i<n;i++)ti[i]=b[i],c[i]=a[i];
    fly(c,1),fly(ti,1);
    for(ll i=0;i<n;i++)c[i]*=ti[i];
    fly(c,-1);
}
cd ge[70001],now[70001],tmp[70001];
ll v[100005];
int main()
{
    ll i,k,j,ans=0,all,maxc=0;
    scanf("%lld",&all);
    s=(ll)sqrt(all)*3;
    for(i=1;i<=all;i++)scanf("%lld",&v[i]),maxc=max(maxc,v[i]);
    for(n=1;n<=maxc*2;n<<=1);
    ll lb=(all-1)/s+1;
    for(i=1;i<=all;i++)now[v[i]]+=1;
    for(i=1;i<=lb;i++){
        for(j=(i-1)*s+1;j<=i*s&&j<=all;j++)
            now[v[j]]-=1;
        FFT(now,ge,tmp);
        for(j=(i-1)*s+1;j<=i*s&&j<=all;j++)
            for(k=j+1;k<=i*s&&k<=all;k++)
                if(2*v[j]-v[k]>=0)ans+=gg(ge[2*v[j]-v[k]]);
        for(j=(i-1)*s+1;j<=i*s&&j<=all;j++)
            ans+=gg(tmp[v[j]*2]);
        for(j=(i-1)*s+1;j<=i*s&&j<=all;j++)
            ge[v[j]]+=1;
    }
    for(i=0;i<=maxc;i++)now[i]=0;
    for(i=all;i;--i){
        ll nb=(i-1)/s+1;
        for(j=i-1;j>=(nb-1)*s+1;--j)
            if(2*v[i]-v[j]>=0)ans+=gg(now[2*v[i]-v[j]]);
        now[v[i]]+=1;
    }
    printf("%lld",ans);
    return 0;
}