首页 > 代码库 > 【点分治】【FFT】Gym - 101234D - Forest Game

【点分治】【FFT】Gym - 101234D - Forest Game

存个求树上每种长度(长度定义为路径上点数)的路径条数的模板:num数组中除了长度为1的以外,都算了2次。

不造为啥FFT数组要开八倍。

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<iostream>
#include<cmath>
using namespace std;
typedef long long ll;
const ll MOD=1000000007ll;
ll ans;
#define MAXN 200010
ll num[MAXN];
const double PI = acos(-1.0);
struct Complex{
    double real,image;
    Complex(double _real,double _image){
        real=_real;
        image=_image;
    }
    Complex(){}
};
Complex operator + (const Complex &c1,const Complex &c2){
    return Complex(c1.real+c2.real,c1.image+c2.image);
}
Complex operator - (const Complex &c1,const Complex &c2){
    return Complex(c1.real-c2.real,c1.image-c2.image);
}
Complex operator * (const Complex &c1,const Complex &c2){
    return Complex(c1.real*c2.real-c1.image*c2.image,c1.real*c2.image+c1.image*c2.real);
}
int rev(int id,int len){
    int ret=0;
    for(int i=0;(1<<i)<len;++i){
        ret<<=1;
        if(id&(1<<i)){
            ret|=1;
        }
    }
    return ret;
}
Complex tmp[800010];
void IterativeFFT(Complex A[],int len, int DFT){
    for(int i=0;i<len;++i){
        tmp[rev(i,len)]=A[i];
    }
    for(int i=0;i<len;++i){
        A[i]=tmp[i];
    }
    for(int s=1;(1<<s)<=len;++s){
        int m=(1<<s);
        Complex wm=Complex(cos(DFT*2*PI/m),sin(DFT*2*PI/m));
        for(int k=0;k<len;k+=m){
            Complex w=Complex(1,0);
            for(int j=0;j<(m>>1);++j){
                Complex t=w*A[k+j+(m>>1)];
                Complex u=A[k+j];
                A[k+j]=u+t;
                A[k+j+(m>>1)]=u-t;
                w=w*wm;
            }
        }
    }
    if(DFT==-1){
        for(int i=0;i<len;++i){
            A[i].real/=len;
            A[i].image/=len;
        }
    }
}
typedef pair<int,int> Point;
int n;
int v[MAXN<<1],w[MAXN<<1],first[MAXN],__next[MAXN<<1],en;
void AddEdge(const int &U,const int &V,const int &W)
{
    v[++en]=V;
    w[en]=W;
    __next[en]=first[U];
    first[U]=en;
}
bool centroid[MAXN];
int size[MAXN];
int calc_sizes(int U,int Fa)
{
    int res=1;
    for(int i=first[U];i;i=__next[i])
      if(v[i]!=Fa&&(!centroid[v[i]]))
        res+=calc_sizes(v[i],U);
    return size[U]=res;
}
Point calc_centroid(int U,int Fa,int nn)
{
    Point res=make_pair(2147483647,-1);
    int sum=1,maxv=0;
    for(int i=first[U];i;i=__next[i])
      if(v[i]!=Fa&&(!centroid[v[i]]))
        {
          res=min(res,calc_centroid(v[i],U,nn));
          maxv=max(maxv,size[v[i]]);
          sum+=size[v[i]];
        }
    maxv=max(maxv,nn-sum);
    res=min(res,make_pair(maxv,U));
    return res;
}
int td[MAXN],en2,ds[MAXN],en3;
void calc_dis(int U,int Fa,int d)
{
    td[en2++]=d;
    for(int i=first[U];i;i=__next[i])
      if(v[i]!=Fa&&(!centroid[v[i]]))
        calc_dis(v[i],U,d+w[i]);
}
Complex fft[800100];
void calc_pairs(int dis[],int En,int op)
{
    int lim=0;
    for(int i=0;i<En;++i){
        lim=max(lim,dis[i]);
    }
    ++lim;
    int len;
    for(int i=0;;++i){
        if((1<<(i-1))>=lim){
            len=(1<<i);
            break;
        }
    }
    for(int i=0;i<len;++i){
        fft[i]=Complex(0,0);
    }
    for(int i=0;i<En;++i){
        fft[dis[i]].real+=1.0;
    }
    IterativeFFT(fft,len,1);
    for(int i=0;i<len;++i){
        fft[i]=fft[i]*fft[i];
    }
    IterativeFFT(fft,len,-1);
    for(int i=1;i<len;++i){
        num[i+1]+=((ll)(fft[i].real+0.5)*(ll)op);
    }
    for(int i=0;i<En;++i){
        num[dis[i]<<1|1]+=(ll)(-op);
    }
}
void solve(int U)
{
    calc_sizes(U,-1);
    int s=calc_centroid(U,-1,size[U]).second;
    centroid[s]=1;
    for(int i=first[s];i;i=__next[i])
      if(!centroid[v[i]])
        solve(v[i]);
    en3=0; ds[en3++]=0;
    for(int i=first[s];i;i=__next[i])
      if(!centroid[v[i]])
        {
          en2=0; calc_dis(v[i],s,w[i]);
          calc_pairs(td,en2,-1);
          memcpy(ds+en3,td,en2*sizeof(int)); en3+=en2;
        }
    calc_pairs(ds,en3,1);
    centroid[s]=0;
}
ll Quick_Pow(ll a,ll p){
	if(!p){
		return 1;
	}
	ll res=Quick_Pow(a,p>>1);
	res=res*res%MOD;
	if((p&1ll)==1ll){
		res=(a%MOD*res)%MOD;
	}
	return res;
}
int main()
{
    int a,b;
    scanf("%d",&n);
    ll njc=1;
    for(int i=1;i<=n;++i){
    	njc=(njc*(ll)i)%MOD;
    }
    for(int i=1;i<n;++i)
      {
        scanf("%d%d",&a,&b);
        AddEdge(a,b,1);
        AddEdge(b,a,1);
      }
    solve(1);
    num[1]=n;
    for(int i=1;i<=n;++i){
    	ans=(ans+((((njc*num[i])%MOD)*Quick_Pow((ll)i,MOD-2ll))%MOD))%MOD;
    }
    cout<<ans<<endl;
    return 0;
}

【点分治】【FFT】Gym - 101234D - Forest Game