codechef Prime Distance On Tree（树分治+FFT）

题目链接：http://www.codechef.com/problems/PRIMEDST/

题意：给出一棵树，边长度都是1。每次任意取出两个点(u,v)，他们之间的长度为素数的概率为多大？

树分治，对于每个根出发记录边的长度出现几次，然后每次求卷积，用素数表查一下即可添加答案。

 #include<algorithm>

 #include<cstdio>

 #include<cmath>

 #include<cstring>

 #include<iostream>

 #include<complex>

 #define ll long long

 #define N 150005

 typedef std::complex<double> cd;

 int mark[],p[];

 int son[],F[],first[],next[],tot,go[];

 int lc,n,sum,dis[],root,vis[],mxdis;

 ll a[],b[],c[],A[],B[];

 ll ans1=,ans2=;

 int read(){

     char ch=getchar();int t=,f=;

     while (ch<''||ch>''){

         if (ch=='-') f=-;

         ch=getchar();

     }

     while (''<=ch&&ch<=''){

         t=t*+ch-'';

         ch=getchar();

     }

     return t*f;

 }

 int bitrev(int t,int n){

     int res=;

     for (int i=;i<n;i++) res|=((t>>(n-i-))&)<<i;//括号要多加

     return res;

 }

 void fft(cd *a,int n,int rev){

     int len=<<n;

     static cd y[N];double Pi=acos(-);

     for (int i=;i<len;i++) y[i]=a[bitrev(i,n)];

     for (int d=;d<len;d<<=){

         cd wn(exp(cd(,Pi*rev/d)));

         for (int k=;k<len;k+=*d){

             cd w(,);

             for (int i=k;i<k+d;i++,w*=wn){

                 cd u=y[i],v=w*y[i+d];

                 y[i]=u+v;

                 y[i+d]=u-v;

             }

         }

     }

     if (rev==-)

     for (int i=;i<len;i++) y[i]/=len;

     for (int i=;i<len;i++) a[i]=y[i];

 }

 void mul(ll *a,int la,ll *b,int lb,ll *c,int &lc){

     int len=,n=;

     static cd t1[N],t2[N];

     for (;len<la*||len<lb*;len<<=,n++);

     for (int i=;i<la;i++) t1[i]=cd(a[i],);

     for (int i=;i<lb;i++) t2[i]=cd(b[i],);

     for (int i=la;i<len;i++) t1[i]=cd(,);

     for (int i=lb;i<len;i++) t2[i]=cd(,);

     fft(t1,n,);fft(t2,n,);

     for (int i=;i<len;i++) t1[i]*=t2[i];

     fft(t1,n,-);

     for (int i=;i<len;i++) c[i]=(ll)(t1[i].real()+0.5);

     lc=len;

 }

 void insert(int x,int y){

     tot++;

     go[tot]=y;

     next[tot]=first[x];

     first[x]=tot;

 }

 void add(int x,int y){

     insert(x,y);

     insert(y,x);

 }

 void prework(){

     mark[]=mark[]=;

     for (int i=;i<=;i++){

         if (!mark[i]){

             p[++p[]]=i;

         }

         for (int j=;j<=p[]&&p[j]*i<=;j++){

             mark[p[j]*i]=;

             if (i%p[j]==) break;

         }

     }

 }

 void findroot(int x,int fa){

     son[x]=;

     F[x]=;

     for (int i=first[x];i;i=next[i]){

         int pur=go[i];

         if (vis[pur]||pur==fa) continue;

         findroot(pur,x);

         son[x]+=son[pur];

         F[x]=std::max(F[x],son[pur]);

     }

     F[x]=std::max(F[x],sum-son[x]);

     if (F[x]<F[root]) root=x;

 }

 void pre(int x,int fa){

     dis[x]=dis[fa]+;

     mxdis=std::max(mxdis,dis[x]);

     for (int i=first[x];i;i=next[i]){

         int pur=go[i];

         if (pur==fa||vis[pur]) continue;

         pre(pur,x);

     }

 }

 void dfs(int x,int fa){

     b[dis[x]]++;

     for (int i=first[x];i;i=next[i]){

         int pur=go[i];

         if (pur==fa||vis[pur]) continue;

         dis[pur]=dis[x]+;

         dfs(pur,x);

     }

 }

 void work(int x){

     dis[]=-;mxdis=;

     int all=sum;

     pre(x,);

     for (int i=;i<=mxdis+;i++) a[i]=b[i]=;

     a[]=;

     for (int i=first[x];i;i=next[i]){

         int pur=go[i];

         if (vis[pur]) continue;

         dis[pur]=;

         for (int j=;j<=mxdis;j++) b[j]=;

         dfs(pur,x);

         for (int j=;j<=mxdis;j++) A[j]=a[j],B[j]=b[j];

         mul(A,mxdis+,B,mxdis+,c,lc);

         int lim=std::min(mxdis*,);

         for (int i=;i<=p[]&&p[i]<=lim;i++){

           ans1+=c[p[i]];

         }

         for (int i=;i<=mxdis;i++)

          a[i]+=b[i];

     }

     vis[x]=;

     for (int i=first[x];i;i=next[i]){

         int pur=go[i];

         if (vis[pur]) continue;

         if (son[pur]>son[x]) sum=all-son[x];

         else sum=son[pur];

         root=;

         findroot(pur,x);

         work(root);

     }

 }

 int main(){

     prework();

     while (scanf("%d",&n)!=EOF){

      if (n==) break;

      tot=;

      for (int i=;i<=n;i++) first[i]=vis[i]=;

      for (int i=;i<n;i++){

         int x,y;

         x=read();y=read();

         add(x,y);

      }

      root=;sum=n;

      F[]=;

      ans1=ans2=;

      findroot(,);

      work(root);

      ans2=(ll)(((ll)n)*((ll)(n-))/);

      double tmp=((double)(ans1))/((double)(ans2));

      printf("%.8lf\n",tmp);

     }

 }