codechef Prime Distance On Tree（树分治+FFT）
题目链接：http://www.codechef.com/problems/PRIMEDST/
题意：给出一棵树，边长度都是1。每次任意取出两个点(u,v)，他们之间的长度为素数的概率为多大？
思路：树分治。对于u的所有孩子两个两个处理。计算到u的距离，然后用FFT求卷积。枚举素数。
struct node
{
double x,y;
node(double _x=0.0,double _y=0.0)
{
x=_x;
y=_y;
}
node operator+(node a)
{
return node(x+a.x,y+a.y);
}
node operator-(node a)
{
return node(x-a.x,y-a.y);
}
node operator*(node a)
{
return node(x*a.x-y*a.y,x*a.y+y*a.x);
}
};
node A[N];
int L;
int reverse(int x)
{
int ans=0,i;
FOR0(i,L) if(x&(1<<i)) ans|=1<<(L-1-i);
return ans;
}
void bitReverseCopy(node a[],int n)
{
int i;
FOR0(i,n) A[i]=a[i];
FOR0(i,n)
{
a[reverse(i)]=A[i];
}
}
void fft(node a[],int n,int on)
{
bitReverseCopy(a,n);
int len,i,j,k;
node x,y,u,t;
for(len=2;len<=n;len<<=1)
{
x=node(cos(-on*2*PI/len),sin(-on*2*PI/len));
for(j=0;j<n;j+=len)
{
y=node(1,0);
for(k=j;k<j+len/2;k++)
{
u=a[k];
t=y*a[k+len/2];
a[k]=u+t;
a[k+len/2]=u-t;
y=y*x;
}
}
}
if(on==-1)
{
FOR0(i,n) a[i].x/=n;
}
}
int prime[N],tag[N],cnt;
void init()
{
tag[0]=1;
tag[1]=1;
int i,j;
for(i=2;i<N;i++) if(!tag[i])
{
prime[cnt++]=i;
for(j=i+i;j<N;j+=i) tag[j]=1;
}
}
vector<int> g[N];
int n,visit[N];
i64 ans;
int f[N],f1[N];
vector<int> V;
void dfs(int u,int pre)
{
f[u]=0;
f1[u]=1;
int i,v;
FOR0(i,SZ(g[u]))
{
v=g[u][i];
if(v==pre||visit[v]) continue;
dfs(v,u);
f1[u]+=f1[v];
upMax(f[u],f1[v]);
}
V.pb(u);
}
int getRoot(int u)
{
V.clear(); dfs(u,0);
int ans=u,p=SZ(V),i,temp,v;
FOR0(i,SZ(V))
{
v=V[i];
temp=max(SZ(V)-f[v],f[v]);
if(temp<p) p=temp,ans=v;
}
return ans;
}
int d[N];
int MaxDis;
node a[N],b[N];
void DFS1(int u,int pre,vector<int> &V)
{
d[u]=d[pre]+1; upMax(MaxDis,d[u]);
V.pb(u);
int i,v;
FOR0(i,SZ(g[u]))
{
v=g[u][i];
if(v==pre||visit[v]) continue;
DFS1(v,u,V);
}
}
vector<int> V1[N];
int P[N],Q[N];
int M;
void deal()
{
int i;
for(i=0;i<=MaxDis;i++) a[i]=node(P[i],0),b[i]=node(Q[i],0);
while(i<M) a[i]=node(0,0),b[i]=node(0,0),i++;
fft(a,M,1);
fft(b,M,1);
for(i=0;i<M;i++) a[i]=a[i]*b[i];
fft(a,M,-1);
i64 temp;
for(i=0;i<cnt&&prime[i]<M;i++)
{
temp=(i64)(a[prime[i]].x+0.5);
ans+=temp;
}
}
void DFS(int u)
{
u=getRoot(u);
MaxDis=0; d[u]=0;
int i,j,v,sonNum=0;
FOR0(i,SZ(g[u]))
{
v=g[u][i];
if(visit[v]) continue;
V1[sonNum].clear();
DFS1(v,u,V1[sonNum]);
sonNum++;
}
M=1; L=0;
while(M<=MaxDis+MaxDis) M<<=1,L++;
for(i=0;i<=MaxDis;i++) P[i]=Q[i]=0;
FOR0(i,sonNum)
{
FOR0(j,SZ(V1[i]))
{
v=V1[i][j];
Q[d[v]]++;
if(!tag[d[v]]) ans++;
}
deal();
for(j=0;j<=MaxDis;j++) P[j]+=Q[j],Q[j]=0;
}
visit[u]=1;
FOR0(i,SZ(g[u]))
{
v=g[u][i];
if(!visit[v]) DFS(v);
}
}
int main()
{
init();
Rush(n)
{
int i;
FOR1(i,n) g[i].clear();
FOR1(i,n-1)
{
int u,v;
RD(u,v);
g[u].pb(v); g[v].pb(u);
}
ans=0; DFS(1);
printf("%.8lf\n",ans/(1.0*n*(n-1)/2));
}
}