Prime Distance On Tree
https://vjudge.net/problem/149908/origin
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.
题意：给你一颗树有n个节点，每条边的长度都是1，现从中任取两点问他们之间的距离是素数的概率
思路：首先树上距离肯定考虑点分治，当已经求出各个点到重心的距离后，考虑如何将两条路径合并
因为路径长度都为1，故对于某一层而言（点分治有log层）最大的路径长度之和不会超过点数n，故将可以用多项式将各个路径长度出现的次数表示出来，一层只有n项，然后自乘即可算出任意两条路径合并后各个长度出现的次数，但是一条路径自己也会和自己合并，故需要减去，将所有长度统计完后将素数长度的相加即可总复杂度
#include<bits/stdc++.h>
#define MAXN 50010
#define INF 0x3f3f3f3f
#define ll long long
#define ri register int
using namespace std;
int head[MAXN],tot;
struct edge
{int v,nxt;
}edg[MAXN << 1];
inline void addedg(int u,int v)
{edg[tot].v = v;edg[tot].nxt = head[u];head[u] = tot++;
}
int prime[MAXN],flag[MAXN],num;
inline void euler(int n)
{for(int i = 2;i <= n;++i){if(!flag[i])prime[++num] = i;for(int j = 1;j <= num && i * prime[j] <= n;++j){flag[i*prime[j]] = 1;if(i % prime[j] == 0)break;}}
}
const double pi=acos(-1.0);
struct cpx
{double r, i;inline void operator +=(const cpx &b){ r+=b.r, i+=b.i;}inline cpx operator +(const cpx &b)const{ return (cpx){r+b.r, i+b.i};}inline cpx operator -(const cpx &b)const{ return (cpx){r-b.r, i-b.i};}inline cpx operator *(const cpx &b)const{ return (cpx){r*b.r-i*b.i, r*b.i+i*b.r};}inline cpx operator *(const double b)const{ return (cpx){r*b, i*b};}inline cpx operator ~()const{return (cpx){r, -i};}
}a[MAXN<<2], w[MAXN<<2];
ll ans[MAXN << 2];
inline void DFT_(cpx *f, int n)
{for(ri i=0, j=0; i<n; ++i){if(i>j) swap(f[i], f[j]);for(ri k=n>>1; (j^=k)<k; k>>=1);}for(ri i=1; i<n; i<<=1) for(ri j=0; j<n; j+=i<<1)for(ri k=j; k<j+i; ++k){cpx t=w[i+k-j]*f[k+i];f[k+i]=f[k]-t, f[k]+=t;}
}
inline void DFT(cpx *f, int n)
{if(n==1) return;n>>=1;static cpx a[MAXN<<1];for(ri i=0; i<n; ++i) a[i]=(cpx){f[i<<1].r, f[i<<1|1].r};DFT_(a, n);for(ri i=0; i<n; ++i){cpx q=~a[(n-i)&(n-1)], x=(a[i]+q)*0.5, y=(a[i]-q)*(cpx){0, -0.5}, t=y*w[n+i];f[i]=x+t, f[n+i]=x-t;}
}
inline void IDFT(cpx *f, int n)
{if(n==1) return;reverse(f+1, f+n), n>>=1;static cpx a[MAXN<<1];for(ri i=0; i<n; ++i)a[i]=(f[i]+f[i+n])*0.5 + (f[i]-f[i+n])*(cpx){0, 0.5}*w[n+i];DFT_(a, n);double k=1.0/n;for(ri i=0; i<n; ++i) f[i<<1]=(cpx){a[i].r*k, 0}, f[i<<1|1]=(cpx){a[i].i*k, 0};
}int n,root,ms,mson[MAXN],sz[MAXN],Size;
bool vis[MAXN];
//root用于标记重心,ms表示树的重心的最大子树的大小,mson[i]记录以i为根最大子树的大小
//sz[i]记录以i为根子树的大小,Size表示当前整棵树的大小,vis[i]表示当前节点是否被分治过
void getroot(int u,int f)//获得重心
{sz[u] = 1,mson[u] = 0;int v;for(int i = head[u];i != -1;i = edg[i].nxt){v = edg[i].v;if(vis[v] || v == f) continue;//剔除已经被分治过的点getroot(v,u);sz[u] += sz[v];if(sz[v] > mson[u]) mson[u] = sz[v];}if(Size - sz[u] > mson[u]) mson[u] = Size-sz[u];//把u看作根节点时u的父亲那一部分也算作子树if(ms > mson[u]) ms = mson[u],root = u;//更新重心
}
int dis[MAXN],cnt;//dis记录所有节点到重心的距离
int mx;
void getdis(int u,int f,int d)//获得到目标点的距离
{dis[++cnt] = d;mx = max(mx,d);int v;for(int i = head[u];i != -1;i = edg[i].nxt){v = edg[i].v;if(vis[v] || v == f) continue;getdis(v,u,d + 1);}
}
void cal(int u,int d,int tp)//u表示getdis的起点,d表示u到目标点的距离,tp表示这一次统计出来的答案是合理的还是不合理的
{cnt = mx = 0;getdis(u,0,d);//算出树中的点到目标点的距离int len = 1;while(len <= mx*2) len <<= 1;for(int i = 0;i < len;++i)a[i] = (cpx){0,0};for(int i = 1;i <= cnt;++i)a[dis[i]].r += 1;for(ri i=1; i<len; i<<=1){w[i]=(cpx){1, 0};for(ri j=1; j<i; ++j)w[i+j]=((j&31)==1?(cpx){cos(pi*j/i), sin(pi*j/i)}:w[i+j-1]*w[i+1]);}DFT(a,len);for(ri i = 0;i < len;++i)a[i] = a[i] * a[i];IDFT(a,len);for(int i = 0;i < len;++i)ans[i] += (ll)(a[i].r+0.5)*tp;for(int i = 1;i <= cnt;++i)ans[dis[i]+dis[i]] -= tp;
}
void solve(int u,int ssize)//ssize是当前这棵子树的大小
{vis[u] = true;//代码保证每次进来的u都必定是当前这棵树的重心,我们将vis[u]标记为true,表示u点被分治过cal(u,0,1);//计算这棵树以u为重心的所有组合,但包括了共用同一条边的情况int v;for(int i = head[u];i != -1;i = edg[i].nxt){v = edg[i].v;if(vis[v]) continue;cal(v,1,-1);//将共用一条边的不合法情况去除ms = INF;//记得每次都要初始化Size = sz[v] < sz[u]?sz[v]:(ssize-sz[u]);//因为v实际上可能是u的父亲,故sz需相减getroot(v,v);//求出以v为根节点的子树重心solve(root,Size);}
}
int mlen;
inline void init()
{tot = 0,ms = INF,Size = n;memset(head,-1,sizeof(int)*(n+1));memset(vis,false,sizeof(bool)*(n+1));mlen = 1;memset(ans,0,sizeof(ll)*(4*n+1));while(mlen < 2*n) mlen <<= 1;
}
int main()
{euler(MAXN-10);while(~scanf("%d",&n)){init();int u,v;for(int i = 1;i < n;++i){scanf("%d%d",&u,&v);addedg(u,v),addedg(v,u);}getroot(1,1);solve(root,Size);
// IDFT(ans,mlen);ll res = 0;for(int i = 1;i <= num && prime[i] < mlen;++i)res += ans[prime[i]];printf("%.6f\n",1.0*res/n/(n-1));}return 0;
}