Codechef PRIMEDST【点分治+FFT】

dzy posted @ 2014年3月21日 11:03 in codechef with tags FFT 点分治 codechef , 2572 阅读

原题传送门

题意:询问在一棵树上任选两点[tex]a,b[/tex]满足[tex]dist(a,b)[/tex]是质数的概率,边权均为1。

 

我们考虑统计树上点对路径为k的点对个数cnt[k]。

显然就是点分。

设一颗子树的dis[k]表示从根下来的路径为k的节点数。

设子树的深度为L,子树都按L从小到大排序,再合并信息,这样复杂度有保证。

在搞第i棵子树的时候,设前i-1棵子树的dis数组为a,第i棵子树的dis数组为b。

一开始只要a[0]=1即可。

统计时,[tex]\text{cnt}[k]=\sum_{i=0}^k \text{a}[i]\text{b}[k-i][/tex]。

然后把b数组加进a数组即可。

暴力统计显然还不如不点分直接暴力= =。

使用FFT,复杂度为[tex]\text{O}(n\text{log}^{2}n)[/tex]。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;
const int maxn=50010,maxm=100020;
const double pi=3.1415926535897932384626;
int p[maxm],n1[maxm],h[maxn],sz[maxn],vis[maxn]={0},ee=0,n;
int size,root,mincnt,maxdis,d[maxn];
long long cnt[maxn],allcnt[maxn],f[maxn],ans;
int isprime[maxn]={0},prime[maxn],pcnt=0,flag[maxn]={0},tmp[100];
#define pr pair<int,int>
#define X first
#define Y second
#define mp make_pair
pr son[maxn];
struct FFT{
	struct cpx{
		double a,b;
		cpx operator + (const cpx&o)const{
			return (cpx){a+o.a,b+o.b};
		}
		cpx operator - (const cpx&o)const{
			return (cpx){a-o.a,b-o.b};
		}
		cpx operator * (const cpx&o)const{
			return (cpx){a*o.a-b*o.b,b*o.a+a*o.b};
		}
	}c[maxn*4],d[maxn*4];
	int len;
	int rev(int x){
		tmp[0]=0;
		for(;x;x>>=1)	tmp[++tmp[0]]=x&1;
		int t=0;
		for(int i=1;i<=tmp[0];i++)	if(tmp[i])	t+=1<<(len-i);
		return t;
	}
	void fft(cpx *A,int n,int fl=1){
		for(int i=0;i<n;i++){
			int t=rev(i);
			if(i<t)	swap(A[i],A[t]);
		}
		for(int m=2;m<=n;m<<=1){
			cpx w=(cpx){cos(fl*2*pi/(double)m),sin(fl*2*pi/(double)m)};
			for(int k=0;k<n;k+=m){
				cpx cur=(cpx){1,0};
				for(int j=k;j<k+(m>>1);j++,cur=cur*w){
					cpx t=cur*A[j+(m>>1)],u=A[j];
					A[j]=u+t;
					A[j+(m>>1)]=u-t;
				}
			}
		}
	}
	long long C[maxn*4],D[maxn*4],ans[maxn*4];
	void merge(int L){
		len=0;
		while((1<<len)<2*L)	len++;
		for(int i=0;i<(1<<len);i++){
			if(i<=L){
				c[i]=(cpx){(double)C[i],0};
				d[i]=(cpx){(double)D[i],0};
			}else{
				c[i]=(cpx){0,0};
				d[i]=(cpx){0,0};
			}
		}
		fft(c,1<<len);
		fft(d,1<<len);
		for(int i=0;i<(1<<len);i++)	c[i]=c[i]*d[i];
		fft(c,1<<len,-1);
		for(int i=0;i<(1<<len);i++)	ans[i]=(long long)((double)c[i].a/(double)(1<<len)+0.5);
	}
}F;
void pre(int n){
	for(int i=2;i<=n;i++){
		if(!flag[i])	isprime[i]=1,prime[++pcnt]=i;
		for(int j=1;j<=pcnt && i*prime[j]<=n;j++){
			flag[i*prime[j]]=1;
			if(i%prime[j]==0)	break;
		}
	}
}
void ae(int x,int y){
	p[ee]=y;	n1[ee]=h[x];	h[x]=ee++;
	p[ee]=x;	n1[ee]=h[y];	h[y]=ee++;
}
void findroot(int u,int fa){
	sz[u]=1;
	int mx=0;
	for(int i=h[u];~i;i=n1[i]){
		if(vis[p[i]] || p[i]==fa)	continue;
		findroot(p[i],u);
		sz[u]+=sz[p[i]];
		mx=max(sz[p[i]],mx);
	}
	mx=max(mx,size-sz[u]);
	if(mx<mincnt){
		mincnt=mx;
		root=u;
	}
}
void dfs(int u,int fa){
	cnt[d[u]]++;
	maxdis=max(maxdis,d[u]);
	for(int i=h[u];~i;i=n1[i]){
		if(p[i]==fa || vis[p[i]])	continue;
		d[p[i]]=d[u]+1;
		dfs(p[i],u);
	}
}
void merge(int l){
	for(int i=0;i<=l;i++)	F.C[i]=allcnt[i],F.D[i]=cnt[i];
	F.merge(l);
	for(int i=1;i<=l;i++)	allcnt[i]+=cnt[i];
	for(int i=1;i<=l;i++)	if(isprime[i])	ans+=F.ans[i];
}
void solve(int u){
	int scnt=0;
	for(int i=h[u];~i;i=n1[i]){
		if(vis[p[i]])	continue;
		d[p[i]]=1;maxdis=0;
		dfs(p[i],u);
		son[++scnt]=mp(maxdis,p[i]);
	}
	sort(son+1,son+scnt+1);
	memset(allcnt,0,sizeof(allcnt));
	allcnt[0]=1;
	son[0].X=1;
	for(int i=1;i<=scnt;i++){
		memset(cnt,0,sizeof(cnt));
		d[son[i].Y]=1;
		dfs(son[i].Y,u);
		merge(son[i-1].X+son[i].X);
	}
	vis[u]=1;
	findroot(u,0);
	for(int i=h[u];~i;i=n1[i]){
		if(vis[p[i]])	continue;
		mincnt=size=sz[p[i]];
		findroot(p[i],u);
		solve(root);
	}
}
int main(){
	scanf("%d",&n);
	pre(n);
	int x,y;
	memset(h,-1,sizeof(h));
	for(int i=1;i<n;i++){
		scanf("%d%d",&x,&y);
		ae(x,y);
	}
	mincnt=size=n;
	findroot(1,0);
	solve(root);
	long long all=1ll*n*(n-1)/2;
	printf("%.8lf\n",(double)ans/(double)all);
	return 0;
}

登录 *


loading captcha image...
(输入验证码)
or Ctrl+Enter