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));
    }
}