CodeChef PRIMEDST Prime Distance On Tree

https://www.codechef.com/problems/PRIMEDST

5e4以下的质数有5000多个，nlogn*5000直接没了，所以不能对于每一个点枚举所有质数，即使加了剪枝也不行

那么我们想到要计算本身和本身相加的和 的方案数，而且本题的边长都是1，所以质数的范围也与当前分治的子树的点数有关，最大的dep为mxdep，那么寻找的质数就不用超过2*mxdep了，所以可以fft

由于fft的长度与分治 的子树的点数有关，且复杂度是lenlognlen的，而点分治的总点数只有nlogn，所以复杂度是n*logn*logn的

#include<bits/stdc++.h>
using namespace std;

const int maxl=5e4+10;
const double pi=acos(-1.0);

int n,cnt,top,len,mxdep;
long long ans;
int a[maxl],p[maxl],s[maxl],dis[maxl],num[maxl*4],ehead[maxl];
bool no[maxl],vis[maxl];
long long res[maxl*4];
struct ed
{
	int to,nxt;
}e[maxl<<1];
struct centertree
{
	int ans,n,mi;
	int son[maxl];
	inline void dfs(int u,int fa)
	{
		int v,mx=0;son[u]=1;
		for(int i=ehead[u];i;i=e[i].nxt)
		{
			v=e[i].to;
			if(vis[v] || v==fa) continue;
			dfs(v,u);
			son[u]+=son[v];
			mx=max(mx,son[v]);
		}
		mx=max(mx,n-son[u]);
		if(mx<mi)
			ans=u,mi=mx;
	}
	inline int getcenter(int u)
	{
		ans=0,mi=maxl;
		dfs(u,0);
		return ans;
	}
}tree;

struct Complex
{
	double r,i;
	Complex(double r1=0.0,double i1=0.0)
	{
		r=r1;i=i1;
	}
	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);
	}
}x1[maxl*4],x2[maxl*4];

inline void change(Complex y[],int len)
{
	int i,j;
	for(i=1,j=len/2;i<len-1;i++)
	{
		if(i<j)
			swap(y[i],y[j]);
		int k=len/2;
		while(j>=k)
			j-=k,k/=2;
		if(j<k)
			j+=k;
	}
}

inline void fft(Complex y[],int len,int on)
{
	change(y,len);
	for(int l=2;l<=len;l<<=1)
	{
		Complex wn(cos(-on*2*pi/l),sin(-on*2*pi/l));
		for(int j=0;j<len;j+=l)
		{
			Complex w(1,0);
			for(int k=j;k<j+l/2;k++)
			{
				Complex u=y[k];
				Complex t=w*y[k+l/2];
				y[k]=u+t;
				y[k+l/2]=u-t;
				w=w*wn;
			}
		}
	}
	if(on==-1)
		for(int i=0;i<len;i++)
			y[i].r/=len;
}

inline void shai()
{
	no[1]=true;int j,t;
	for(int i=2;i<maxl;i++)
	{
		if(!no[i]) p[++p[0]]=i;
		j=1,t=p[1]*i;
		while(t<maxl && j<=p[0])
		{
			no[t]=true;
			if(i%p[j]==0) 
				break;
			t=i*p[++j];
		}
	}
}

inline void add(int u,int v)
{
	e[++cnt].to=v;e[cnt].nxt=ehead[u];ehead[u]=cnt;
}

inline void prework()
{
	cnt=0;
	for(int i=1;i<=n;i++)
		ehead[i]=0,vis[i]=false;
	int u,v;
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&u,&v);
		add(u,v);add(v,u);
	}
}

inline void dfs(int u,int fa)
{
	int v;s[++top]=dis[u];
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(vis[v] || v==fa) continue;
		dis[v]=dis[u]+1;
		dfs(v,u);
	}
}

inline long long calc(int u,int w)
{
	int v;len=0;mxdep=0;
	dis[u]=w;top=0;
	dfs(u,0);
	
	for(int i=1;i<=top;i++)
		num[s[i]]++,mxdep=max(s[i],mxdep);
	int xlen=1;
	while(xlen<=2*mxdep)
		xlen<<=1;
	for(int i=0;i<xlen;i++)
	{
		if(i<=mxdep) x1[i]=x2[i]=Complex(num[i],0);
		else x1[i]=x2[i]=Complex(0,0);
	}
	fft(x1,xlen,1);
	fft(x2,xlen,1);
	for(int i=0;i<xlen;i++)
		x1[i]=x1[i]*x2[i];
	fft(x1,xlen,-1);
	for(int i=0;i<xlen;i++)
		res[i]=(long long)(x1[i].r+0.5);
	for(int i=1;i<=top;i++)
		res[s[i]+s[i]]--;
	for(int i=0;i<xlen;i++)
		res[i]/=2;
	long long sum=0;
	for(int i=1;i<=p[0] && p[i]<=mxdep*2;i++)
		sum+=res[p[i]];
	for(int i=1;i<=top;i++)
		num[s[i]]--;
	return sum;
}

inline void solv(int u)
{
	vis[u]=true;
	ans+=calc(u,0);int v,rt;
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(vis[v]) continue;
		ans-=calc(v,1);
		tree.n=tree.son[v];
		rt=tree.getcenter(v);
		solv(rt);
	}
}

inline void mainwork()
{
	ans=0;tree.n=n;
	int rt=tree.getcenter(1);
	solv(rt);
}

inline void print()
{
	double pp=1.0*ans/(1ll*n*(n-1)/2);
	printf("%.7f\n",pp);
}

int main()
{
	shai();
	while(~scanf("%d",&n))
	{
		prework();
		mainwork();
		print();
	}
	return 0;
}