2019.01.11【CodeChef PRIMEDST】Prime Distance On Tree（点分治）（FFT）

版权声明：转载请声明出处，谢谢配合。 https://blog.csdn.net/zxyoi_dreamer/article/details/86314108

传送门

---

解析：

很裸的点分治，很裸的卷积直接上FFT。

思路：

这个问题，统计路径，显然直接上点分治啊。

每次统计出到分治重心距离为 $d$的有多少个，记录到一个数组 $A$里面，然后我们发现需要将两条路径拼起来再计算出现次数。。。

然后发现这就是自我卷积。。。直接上 $FFT$然后利用预先筛出的质数表来统计答案就行了。

---

UPD:

去其他的博客上逛了一圈，好像有点不对。。。

没有太看懂其他人的题解，感觉好像特殊处理了两个 $1$拼起来成2的情况？

这个不是会在处理子树重复统计的时候会被容斥掉吗。。。

反正我看到的题解不是单独在卷积前处理1就是在卷积后处理2？

---

代码：

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define re register
#define gc get_char
#define pc putchar
#define cs const

namespace IO{
	namespace IOONLY{
		cs int Rlen=1<<16|1;
		char buf[Rlen],*p1,*p2;
	}
	inline char get_char(){
		using namespace IOONLY;
		return (p1==p2)&&(p2=(p1=buf)+fread(buf,1,Rlen,stdin),p1==p2)?EOF:*p1++;
	}
	
	inline int getint(){
		re int num;
		re char c;
		while(!isdigit(c=gc()));num=c^48;
		while(isdigit(c=gc()))num=(num+(num<<2)<<1)+(c^48);
		return num;
	}
}
using namespace IO;

cs int P=50004;
int prime[P],pcnt;
bool mark[P];
inline void linear_sieves(int len=P-4){
	for(int re i=2;i<=len;++i){
		if(!mark[i])prime[++pcnt]=i;
		for(int re j=1;i*prime[j]<=len;++j){
			mark[i*prime[j]]=true;
			if(i%prime[j]==0)break;
		}
	}
}

cs int N=50004;
cs double PI=acos(-1);
struct Complex{
	double x,y;
	Complex(){}
	Complex(cs double _x,cs double _y):x(_x),y(_y){}
	inline friend Complex operator+(cs Complex &a,cs Complex &b){return Complex(a.x+b.x,a.y+b.y);}
	inline friend Complex operator-(cs Complex &a,cs Complex &b){return Complex(a.x-b.x,a.y-b.y);}
	inline friend Complex operator*(cs Complex &a,cs Complex &b){return Complex(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);}
}A[N<<2];
int r[N<<2];

inline void FFT(Complex *A,int len,int typ){
	for(int re i=0;i<len;++i)r[i]=r[i>>1]>>1|((i&1)*(len>>1));
	for(int re i=0;i<len;++i)if(i<r[i])swap(A[i],A[r[i]]);
	for(int re i=1;i<len;i<<=1){
		Complex Wn=Complex(cos(PI/i),typ*sin(PI/i));
		for(int re j=0;j<len;j+=i<<1){
			Complex w=Complex(1,0),x,y;
			for(int k=0;k<i;++k,w=w*Wn){
				x=A[j+k],y=A[i+j+k]*w;
				A[j+k]=x+y;
				A[i+j+k]=x-y;
			}
		}
	}
}

vector<int> edge[N];
inline void addedge(int u,int v){
	edge[u].push_back(v);
	edge[v].push_back(u);
}

bool ban[N];
int siz[N],mx,total,G,n;
ll ans;
void get_siz(int u,int fa){
	siz[u]=1;
	for(int re e=0;e<edge[u].size();++e){
		re int v=edge[u][e];
		if(ban[v])continue;
		if(v^fa){
			get_siz(v,u);
			siz[u]+=siz[v];
		}
	}
}

void find_G(int u,int fa){
	int maxn=total-siz[u];
	for(int re e=0;e<edge[u].size();++e){
		re int v=edge[u][e];
		if(ban[v])continue;
		if(v^fa){
			find_G(v,u);
			maxn=max(maxn,siz[v]);
		}
	}
	if(mx>maxn)mx=maxn,G=u;
}

int mxdep;

void get_dis(int u,int fa,int dep){
	mxdep=max(dep,mxdep);
	++A[dep].x;
	for(int re e=0;e<edge[u].size();++e){
		re int v=edge[u][e];
		if(ban[v]||v==fa)continue;
		get_dis(v,u,dep+1);
	}
}

inline ll calc(int u,int now){
	mxdep=0;get_dis(u,0,now);
	mxdep*=2;
	int len=1;ll ans=0;
	for(;len<=mxdep;len<<=1);
	FFT(A,len,1);
	for(int re i=0;i<len;++i)A[i]=A[i]*A[i];
	FFT(A,len,-1);
	for(int re i=1;i<=pcnt&&prime[i]<=mxdep*2;++i){
		ans+=(A[prime[i]].x/len+0.5);
	}
	for(int re i=0;i<len;++i)A[i]=Complex(0,0);
	return ans;
}

inline void solve_G(int u){
	ban[u]=true;
	ans+=calc(u,0);
	for(int re e=0;e<edge[u].size();++e){
		re int v=edge[u][e];
		if(ban[v])continue;
		ans-=calc(v,1);
	}
	for(int re e=0;e<edge[u].size();++e){
		re int v=edge[u][e];
		if(ban[v])continue;
		get_siz(v,u);
		total=mx=siz[v],G=v;
		find_G(v,u);
		solve_G(G);
	}
}

signed main(){
	n=getint();linear_sieves();
	for(int re i=1;i<n;++i)addedge(getint(),getint());
	get_siz(1,0);
	total=mx=siz[1],G=1;
	find_G(1,0);
	solve_G(G);
	printf("%.10lf",(double)(ans)/n/(n-1));
	return 0;
}