HDU5307 He is Flying

題目鏈接

題目大意

給定n 個非負整數a1 ~an ,對每個0≤S≤∑ai ,輸出所有和爲S 的區間的長度和.

1≤n≤105,∑ai≤5×104.

題解

寫這題真心漲姿勢了…
官方題解構造了這樣一個多項式:

(∑ixsi)(∑x−si−1)−(∑xsi)(∑(i−1)x−si−1).

其中si=∑ij=1aj(1≤i≤n),s0=0 .
然後得到的多項式的xS(1≤S≤∑ai) 的係數即爲所求的S 對應的答案.
一開始看到這個我完全是不知所措的,覺得題解真是用心險惡.
然而這個其實還是好理解的,區間[i,j] 的和爲S ,很容易想到化成sj−si−1=S .而上面的多項式左邊兩項乘起來是這樣:

ixsi⋅x−sj−1=ixsi−sj−1.

右邊的兩項乘起來是這樣:

xsi⋅(j−1)x−sj−1=(j−1)xsi−sj−1.

兩個相減就得到了這麼個東西:

(i−j+1)xsi−sj−1.

於是十分巧妙地把S 體現在了指數上,把要求的區間長度體現在了係數上.
而且這樣也避免了遺漏和重複計算,除了S=0 的情況,因爲S=0 可以由si−si 得到,也可以由兩個相等的前綴和si=sj 通過si−sj 和sj−si 兩次計算得到,這樣會比較麻煩.而實際上直接在原數組ai 上O(n) 掃一遍就可以得出S=0 的答案了,所以可以另外考慮.
所以接下來只要算那個構造出來的多項式就好.這個可以用FFT,但是由於指數會有負的,我們把指數爲負的先加上一個值(比如sn ),然後就可以直接用FFT了.
說起來很簡單,但是實現還是有點小煩的.
還有這題會卡精度所以要用long double.
總的時間複雜度是O(nlgn) .
後來手寫了個複數竟然快了近三倍orz…
說寫這題漲姿勢是因爲調試時發現long double要用%Ld輸出…

代碼

#include<cstdio>
#include<cmath>
#include<cstring>
#include<iostream>
#include<algorithm>
#define debug(x) cout<<#x<<"="<<x<<endl
using namespace std;
typedef long long ll;
typedef long double ld;
const int N=1e5+5,M=(1<<17)+5;
const ld pi=acos(-1.0);
int sum[N],rev[M];
struct C{
    ld real,imag;
    C(ld real=0,ld imag=0):real(real),imag(imag){}
    friend C operator +(C A,C B){
        return C(A.real+B.real,A.imag+B.imag);
    }
    friend C operator -(C A,C B){
        return C(A.real-B.real,A.imag-B.imag);
    }
    friend C operator *(C A,C B){
        return C(A.real*B.real-A.imag*B.imag,A.imag*B.real+A.real*B.imag);
    }
}A[M],B[M],ans[2][M];
void rd(int &res){
    res=0;
    char c;
    while(c=getchar(),c<48);
    do res=(res<<3)+(res<<1)+(c^48);
        while(c=getchar(),c>47);
}
void FFT(C *arr,int n,int flag){
    for(int i=0;i<n;++i)
        if(i<rev[i])
            swap(arr[i],arr[rev[i]]);
    for(int m=2;m<=n;m<<=1){
        C wm(cos(2*pi/m),flag*sin(2*pi/m));
        for(int i=0;i<n;i+=m){
            C w(1,0);
            for(int j=0;j<m>>1;++j,w=w*wm){
                C x=arr[i+j],y=w*arr[i+j+(m>>1)];
                arr[i+j]=x+y;
                arr[i+j+(m>>1)]=x-y;
            }
        }
    }
}
void calc_FFT(int n,bool id){
    int _n=n,m=n<<1,S=0;
    for(n=1;n<=m;n<<=1,++S);
    rev[0]=0;
    for(int i=1;i<n;++i)
        rev[i]=(rev[i>>1]>>1)|((i&1)<<(S-1));
    for(int i=_n+1;i<n;++i)
        A[i]=B[i]=0;
    FFT(A,n,1);
    FFT(B,n,1);
    for(int i=0;i<n;++i)
        ans[id][i]=A[i]*B[i];
    FFT(ans[id],n,-1);
    for(int i=0;i<=m;++i)
        ans[id][i].real/=n;
}
void solve(){
    int n;
    ll ans0=0;
    rd(n);
    sum[0]=0;
    for(int i=1,num,cnt=0;i<=n;++i){
        rd(num);
        sum[i]=sum[i-1]+num;
        if(!num){
            ++cnt;
            ans0+=1ll*cnt*(cnt+1)/2;
        }
        else cnt=0;
    }
    printf("%I64d\n",ans0);
    for(int i=0;i<=sum[n];++i){
        A[i]=B[i]=0;
    }
    for(int i=1;i<=n;++i){
        A[sum[i]].real+=i;
        B[-sum[i-1]+sum[n]].real+=1;
    }
    calc_FFT(sum[n],0);
    for(int i=0;i<=sum[n];++i){
        A[i]=B[i]=0;
    }
    for(int i=1;i<=n;++i){
        A[sum[i]].real+=1;
        B[-sum[i-1]+sum[n]].real+=i-1;
    }
    calc_FFT(sum[n],1);
    for(int i=1;i<=sum[n];++i){
        printf("%I64d\n",(ll)(ans[0][i+sum[n]].real-ans[1][i+sum[n]].real+0.5));
    }
}
int main(){
    int cas;
    rd(cas);
    while(cas--)solve();
    return 0;
}
/*

    Jun.19.16

    Tags:math,FFT
    Submissions:5

    Exe.Time 1107MS
    Exe.Memory 18880K
    Code Len. 2239B


*/