[CODECHEF]Prime Distance On Tree

Prime Distance On Tree

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

Problem description.

You are given a tree. If we select 2 distinct nodes uniformly at random, what’s the probability that the distance between these 2 nodes is a prime number?
Input

The first line contains a number N: the number of nodes in this tree.
The following N-1 lines contain pairs a[i] and b[i], which means there is an edge with length 1 between a[i] and b[i].
Output

Output a real number denote the probability we want.
You’ll get accept if the difference between your answer and standard answer is no more than 10^-6.
Constraints

2 ≤ N ≤ 50,000

The input must be a tree.
Example

Input:
5
1 2
2 3
3 4
4 5

Output:
0.5
Explanation

We have C(5, 2) = 10 choices, and these 5 of them have a prime distance:

1-3, 2-4, 3-5: 2

1-4, 2-5: 3

题意是球树上距离为质数的点对个数。
很明显的对于这种球树上点对个数的问题，需要点分。
因为质数有很多，所以我们不能一个一个的去算，这样就T了。可以先写出每一层计算的式子,先不考虑质数的情况，那么：

$$ans[i]=\sum_{j=0}^{i-1}sum[j]*sum[i-j]$$
ans[i]表示距离为i的点对数，sum[i]表示力当前点距离为i的点的个数。那么很明显这就是一个卷积的形式，所以我们就可以用FFT去优化一下，最后再把质数的提出来就好了。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
using namespace std;
#define LL long long
#define inf 0x7fffffff
const int M=200010;
LL ans;
bool flag[M],use[M];
struct S{int st,en;}aa[M*2];
int n,tot,point[M],next[M*2],prime[M],root,f[M],siz[M],deep[M],dig[M],rev[M],N,L,sum,now_max;
struct T{
    double x,y;
    void prepare(double xx,double yy){x=xx;y=yy;}
    T operator + (const T&xx){return (T){x+xx.x,y+xx.y};}
    T operator - (const T&xx){return (T){x-xx.x,y-xx.y};}
    T operator * (const T&xx){return (T){x*xx.x-y*xx.y,x*xx.y+y*xx.x};}
}a[M],b[M],c[M],d[M];
inline int in(){
    int x=0; char ch=getchar();
    while(ch<'0'||ch>'9') ch=getchar();
    while(ch>='0'&&ch<='9') x=x*10+ch-'0',ch=getchar();
    return x;
}
inline void add(int x,int y){
    tot+=1;next[tot]=point[x];point[x]=tot;
    aa[tot].st=x;aa[tot].en=y;
    tot+=1;next[tot]=point[y];point[y]=tot;
    aa[tot].st=y;aa[tot].en=x;
}
inline void get_prime(){
    int i,j;
    for(i=2;i<=M;++i){
        if(!flag[i]) prime[++prime[0]]=i;
        for(j=1;j<=prime[0]&&i*prime[j]<M;++j){
            flag[i*prime[j]]=true;
            if(i%prime[j]==0) break;
        }
    }
}
inline void FFT(T *a,int f){
    int i,j,k;
    T wn,w,x,y;
    for(i=0;i<N;++i) d[i]=a[rev[i]];
    for(i=0;i<N;++i) a[i]=d[i];
    for(i=2;i<=N;i<<=1){
        wn.prepare(cos(2*M_PI/i),f*sin(2*M_PI/i));
        for(j=0;j<N;j+=i){
            w.prepare(1,0);
            for(k=j;k<j+i/2;++k){
                x=a[k];
                y=w*a[k+i/2];
                a[k]=x+y;
                a[k+i/2]=x-y;
                w=w*wn;
            }
        }
    }
    if(f==-1) for(i=0;i<N;++i) a[i].x/=(double)N;
}
inline void get_root(int x,int last){
    int i;
    siz[x]=1;f[x]=0;
    for(i=point[x];i;i=next[i])
      if(aa[i].en!=last&&use[aa[i].en]){
        get_root(aa[i].en,x);
        siz[x]+=siz[aa[i].en];
        f[x]=max(f[x],siz[aa[i].en]);
      }
    f[x]=max(f[x],sum-siz[x]);
    if(f[x]<f[root]) root=x;
}
inline void get_deep(int x,int last,int now){
    int i;
    deep[now]+=1;
    now_max=max(now_max,now);
    for(i=point[x];i;i=next[i])
      if(use[aa[i].en]&&aa[i].en!=last)
        get_deep(aa[i].en,x,now+1);
}
inline LL calc(int x,int y){
    LL now=0;
    int i,j,len;
    for(i=0;i<=now_max;++i) deep[i]=0;
    now_max=0;get_deep(x,0,y);now_max+=1;
    for(N=1,L=0;N<now_max;N<<=1,L+=1); N<<=1;L+=1;
    for(i=0;i<N;++i) dig[i]=0;
    for(i=0;i<N;++i){
        rev[i]=0;
        for(j=i,len=0;j;j>>=1) dig[len++]=j&1;
        for(j=0;j<L;++j) rev[i]=rev[i]*2+dig[j];
    }
    for(i=0;i<N;++i){
        a[i].prepare(i>=N?0:deep[i],0);
        b[i].prepare(i>=N?0:deep[i],0);
    }
    FFT(a,1);FFT(b,1);
    for(i=0;i<N;++i) c[i]=a[i]*b[i];
    FFT(c,-1);
    for(i=1;i<=prime[0]&&prime[i]<N;++i)
      now+=(LL)(c[prime[i]].x+0.5);
    return now;
}
inline void work(int x){
    int i;
    use[x]=false;
    ans+=calc(x,0);
    for(i=point[x];i;i=next[i])
      if(use[aa[i].en]){
        ans-=calc(aa[i].en,1);
        root=0;sum=siz[aa[i].en];
        get_root(aa[i].en,0);
        work(root);
      }
}
int main(){
    int i,j,x,y;
    n=in();
    for(i=1;i<n;++i){
        x=in();y=in();
        add(x,y);
    }
    get_prime();
    sum=n;root=0;f[0]=inf;
    memset(use,1,sizeof(use));
    get_root(1,0);
    work(root);
    LL all=(LL)n*(LL)(n-1);
    printf("%.9f\n",(double)ans/(double)all);
}