Prime Distance On Tree-树分治+FFT

题目描述

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
Note that 1 is not a prime number.

题目分析

使用树分治后我们能够得到任意点到重心的距离，现在问题转换为如何判断两点之间的距离为素数。

首先我们要使用线性筛处理出一个素数表。

一个简单的想法是我们将Dis数组中不同的点加起来就是不同两点之间的距离。然后再判断这个距离是不是素数，如果是素数则将路径+1。可是这样时间复杂度太高（O(n*n)），因此我们需要使用FFT进行加速，快速得到两路径的和。但是我们需要减去同一条路径走两边的情况。

需要求概率，我们就需要知道事件数是什么。对于这里来讲，我们需要考虑的是选择排列还是组合。由给出的样例我们看出，应该按照组合进行计算。因此对于每个路径和我们都应该/2，这样得到的才是排列数。

需要注意的是每次FFT以后对数组的清零。我就是这里没有注意然后导致疯狂出错，还一直检查不出问题。

还有就是，树上的路径很多，很容易爆long long

AC代码

#include<iostream>
#include<cstring>
#include<cstdio>
#include<climits>
#include<algorithm>
#include<ctime>
#include<cstdlib>
#include<queue>
#include<set>
#include<map>
#include<cmath>using namespace std;const int MAXN=5e5+5;
const double PI=acos(-1.0);
typedef long long ll;
struct edge
{int to,len,last;
}Edge[MAXN<<2]; int Last[MAXN],tot;
int n,kk,SonNum[MAXN],MaxNum[MAXN],Vis[MAXN],Dis[MAXN];
int Prime[MAXN]; bool IsPrime[MAXN]; int prime_num=0;
int root,rootx,dlen,ss;
ll Num[MAXN<<2],MaxLen,len;
ll Res[MAXN<<2];
ll ans;struct complex
{double r,i;complex(double _r=0,double _i=0):r(_r),i(_i){}complex operator +(const complex &b) {return complex(r+b.r,i+b.i);}complex operator -(const complex &b) {return complex(r-b.r,i-b.i);}complex operator *(const complex &b) {return complex(r*b.r-i*b.i,r*b.i+i*b.r);}
}A[MAXN<<2];void change(complex y[],int len)
{int i,j,k;for(i = 1, j = len/2;i < len-1;i++){if(i < j)swap(y[i],y[j]);k = len/2;while( j >= k){j -= k;k /= 2;}if(j < k)j += k;}
}void fft(complex y[],int len,int on)
{change(y,len);for(int h = 2;h <= len;h <<= 1){complex wn(cos(-on*2*PI/h),sin(-on*2*PI/h));for(int j = 0;j < len;j += h){complex w(1,0);for(int k = j;k < j+h/2;k++){complex u = y[k];complex t = w*y[k+h/2];y[k] = u+t;y[k+h/2] = u-t;w = w*wn;}}}if(on == -1)for(int i = 0;i < len;i++)y[i].r /= len;
}void FFT(ll a[],int la,ll b[])//la,lb分别是a,b数组的最高位+1 
{//int len=1; while(len<la+la) len<<=1;for(int i=0;i<la;++i) A[i]=complex(a[i],0);for(int i=la;i<len;++i) A[i]=complex(0,0);fft(A,len,1);for(int i=0;i<len;++i) A[i]=A[i]*A[i];fft(A,len,-1);for(int i=0;i<len;++i) b[i]=(ll)(A[i].r+0.5);
}void CreatPrime()
{IsPrime[0]=IsPrime[1]=true; prime_num=0;for(int i=2;i<MAXN;++i){if(!IsPrime[i])Prime[prime_num++]=i;for(int j=0;j<prime_num && Prime[j]*i<MAXN;j++){IsPrime[Prime[j]*i]=true;if(i%Prime[j]==0) break;}}
}int getint()
{int x=0,sign=1; char c=getchar();while(c<'0' || c>'9'){if(c=='-') sign=-1; c=getchar();}while(c>='0' && c<='9'){x=x*10+c-'0'; c=getchar();}return x*sign;
}void Init()
{//CreatPrime();//for(int i=0;i<=n;++i) Last[i]=0; memset(Last,-1,sizeof(Last));tot=0;ans=0; 
}void Clear()
{for(int i=0;i<=n;++i) Vis[i]=false;
}void AddEdge(int u,int v,int w)
{Edge[++tot].to=v; Edge[tot].len=w; Edge[tot].last=Last[u]; Last[u]=tot;
}void Read()
{n=getint();int u,v;for(int i=1;i<n;i++){u=getint(); v=getint();AddEdge(u,v,1); AddEdge(v,u,1);}
}void GetRoot(int x,int father)
{int v;SonNum[x]=1; MaxNum[x]=1;for(int i=Last[x];~i;i=Edge[i].last){v=Edge[i].to; if(v==father || Vis[v]) continue;GetRoot(v,x);SonNum[x]+=SonNum[v];if(SonNum[v]>MaxNum[x]) MaxNum[x]=SonNum[x];}if(ss-SonNum[x]>MaxNum[x]) MaxNum[x]=ss-SonNum[x];if(rootx>MaxNum[x]) root=x,rootx=MaxNum[x];
}void GetDis(int x,int father,int dis)
{int v;Dis[++dlen]=dis;for(int i=Last[x];~i;i=Edge[i].last){v=Edge[i].to; if(v==father|| Vis[v]) continue;GetDis(v,x,dis+Edge[i].len);}
}ll Count(int x,int dis)
{ll ret=0;for(int i=0;i<=dlen;++i) Dis[i]=0;dlen=0;GetDis(x,0,dis);/*for(int i=1;i<=dlen;++i)for(int j=i+1;j<=dlen;++j){if(!IsPrime[Dis[i]+Dis[j]]) ++ret;}*///memset(Num,0,sizeof(Num));MaxLen=0;for(int i=1;i<=dlen;++i){++Num[Dis[i]]; if(Dis[i]>MaxLen) MaxLen=Dis[i];}len=1; while(len<=2*MaxLen) len<<=1;FFT(Num,MaxLen+1,Res);for(int i=1;i<=dlen;++i){--Res[Dis[i]+Dis[i]];}MaxLen<<=1;for(int i=0;i<=len;++i){Res[i]/=2;}for(int i=0;i<prime_num && Prime[i]<=MaxLen;++i){ret+=Res[Prime[i]];}for(int i=1;i<=dlen;++i){Num[Dis[i]]=0;}return ret;
}void Solve(int x)
{int v;ans+=Count(x,0);Vis[x]=true;for(int i=Last[x];~i;i=Edge[i].last){v=Edge[i].to; if(Vis[v]) continue;ans-=Count(v,Edge[i].len);ss=SonNum[v]; rootx=INT_MAX; root=0;GetRoot(v,x);Solve(root);}
}void Work()
{rootx=INT_MAX; ss=n; root=0;GetRoot(1,0); Solve(root);
}void Write()
{ll tmp=(ll)n*(n-1)/2;//printf("%.f\n",tmp);printf("%.7f\n",(double)ans/tmp);}int main()
{CreatPrime();//while(~scanf("%d",&n))//{Init();Read();Work();Write();Clear();//}return 0;
}