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);
     }
 }