【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;

}