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]表示距离为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);
}