【CodeChef】Prime Distance On Tree

vjudge

给定一棵边长都是\(1\)的树，求有多少条路径长度为质数

树上路径自然是点分治去搞，但是发现要求是长度为质数，总不能对每一个质数都判断一遍吧

自然是不行的，这个东西显然是一个卷积，我们合并的时候显然可以直接大力\(NTT\)

但是需要注意的是我们访问子树的顺序必须是先访问深度小的子树，否则轻松被菊花加长链卡掉

但是\(CodeChef\)数据水啊，就这样直接搞过去了

代码

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<set>
#define re register
#define LL long long
#define max(a,b) ((a)>(b)?(a):(b))
#define min(a,b) ((a)<(b)?(a):(b))
const int maxn=1e5+5;
const int inf=1e9;
inline int read() {
	char c=getchar();int x=0;while(c<'0'||c>'9') c=getchar();
	while(c>='0'&&c<='9') x=(x<<3)+(x<<1)+c-48,c=getchar();return x;
}
struct E{int v,nxt;}e[maxn];
const LL mod=998244353;
const LL G[2]={3,332748118};
int sum[maxn],mx[maxn],vis[maxn],st[maxn];
int now,S,rt,top,n,num,head[maxn],tax[maxn];
int f[maxn],p[maxn>>1];
int rev[262145],len,L;
LL A[262145],B[262145],ans;
inline void add(int x,int y) {e[++num].v=y;e[num].nxt=head[x];head[x]=num;}
inline LL ksm(LL a,int b) {
	LL S=1;
	while(b) {if(b&1) S=S*a%mod;b>>=1;a=a*a%mod;}
	return S;
}
inline void NTT(LL *f,int o)  {
	for(re int i=0;i<len;i++) 
		if(i<rev[i]) std::swap(f[i],f[rev[i]]);
	for(re int i=2;i<=len;i<<=1) {
		int ln=i>>1;LL og1=ksm(G[o],(mod-1)/i);
		for(re int l=0;l<len;l+=i) {
			LL t,og=1;
			for(re int x=l;x<l+ln;++x) {
				t=(og*f[ln+x])%mod;
				f[ln+x]=(f[x]-t+mod)%mod;
				f[x]=(f[x]+t)%mod;
				og=(og*og1)%mod;
			}
		}
	}
	if(!o) return;
	LL inv=ksm(len,mod-2);
	for(re int i=0;i<len;i++) f[i]=(f[i]*inv)%mod;
}
inline void mul() {
	int m=0;
	for(re int i=1;i<=top;i++) m=max(m,st[i]);
	if(!L) {L=m;return;}
	for(re int i=1;i<=top;i++) B[st[i]]++;
	len=1;while(len<m+L+2) len<<=1;
	for(re int i=0;i<len;i++) rev[i]=rev[i>>1]>>1|((i&1)?len>>1:0);
	for(re int i=1;i<=L;i++) A[i]=tax[i];
	NTT(A,0),NTT(B,0);
	for(re int i=0;i<len;i++) A[i]=(A[i]*B[i])%mod;
	NTT(A,1);
	for(re int i=1;i<=p[0]&&p[i]<=L+m;i++) ans+=A[p[i]];
	L=max(L,m);
	for(re int i=0;i<len;i++) A[i]=B[i]=0;
}
void getroot(int x,int fa) {
	sum[x]=1;mx[x]=0;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]||e[i].v==fa) continue;
		getroot(e[i].v,x);
		sum[x]+=sum[e[i].v];
		mx[x]=max(mx[x],sum[e[i].v]);
	}
	mx[x]=max(mx[x],S-sum[x]);
	if(mx[x]<now) now=mx[x],rt=x;
}
void getdis(int x,int fa,int L) {
	st[++top]=L;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]||e[i].v==fa) continue;
		getdis(e[i].v,x,L+1);
	}
}
void dfs(int x) {
	vis[x]=1;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]) continue;
		top=0;getdis(e[i].v,x,1);
		for(re int j=1;j<=top;j++)
		if(!f[st[j]]) ans++;
		mul();
		for(re int j=1;j<=top;j++) tax[st[j]]++;
	}
	for(re int i=1;i<=L;i++) tax[i]=0;
	L=0;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]) continue;
		S=sum[e[i].v];now=inf;
		getroot(e[i].v,x);dfs(rt);
	}
}
int main() {
	n=read();
	for(re int x,y,i=1;i<n;i++)
		x=read(),y=read(),add(x,y),add(y,x);
	f[1]=1;
	for(re int i=2;i<=n;i++) {
		if(!f[i]) p[++p[0]]=i;
		for(re int j=1;j<=p[0]&&p[j]*i<=n;j++) {
			f[p[j]*i]=1;
			if(i%p[j]==0) break;
		}
	}
	S=n,now=inf;getroot(1,0);dfs(rt);
	LL C=(LL)n*(LL)(n-1)/2ll;
	printf("%.6lf\n",(double)ans/(double)C);
	return 0;
}