CodeChef PRIMEDST,点分治,FFT

Prime Distance On Tree

https://vjudge.net/problem/149908/origin
Problem description.
You are given a tree. If we select 2 distinct nodes uniformly at random, what’s the probability that the distance between these 2 nodes is a prime number?

Input
The first line contains a number N: the number of nodes in this tree.
The following N-1 lines contain pairs a[i] and b[i], which means there is an edge with length 1 between a[i] and b[i].

Output
Output a real number denote the probability we want.
You’ll get accept if the difference between your answer and standard answer is no more than 10^-6.

题意：给你一颗树有n个节点，每条边的长度都是1，现从中任取两点问他们之间的距离是素数的概率
思路：首先树上距离肯定考虑点分治，当已经求出各个点到重心的距离后，考虑如何将两条路径合并
因为路径长度都为1，故对于某一层而言（点分治有log层）最大的路径长度之和不会超过点数n，故将可以用多项式将各个路径长度出现的次数表示出来，一层只有n项，然后自乘即可算出任意两条路径合并后各个长度出现的次数，但是一条路径自己也会和自己合并，故需要减去，将所有长度统计完后将素数长度的相加即可总复杂度 $O(n*log^2n)$

#include<bits/stdc++.h>
#define MAXN 50010
#define INF 0x3f3f3f3f
#define ll long long
#define ri register int
using namespace std;
int head[MAXN],tot;
struct edge
{
    int v,nxt;
}edg[MAXN << 1];
inline void addedg(int u,int v)
{
    edg[tot].v = v;
    edg[tot].nxt = head[u];
    head[u] = tot++;
}
int prime[MAXN],flag[MAXN],num;
inline void euler(int n)
{
    for(int i = 2;i <= n;++i)
    {
        if(!flag[i])
            prime[++num] = i;
        for(int j = 1;j <= num && i * prime[j] <= n;++j)
        {
            flag[i*prime[j]] = 1;
            if(i % prime[j] == 0)
                break;
        }
    }
}
const double pi=acos(-1.0);
struct cpx
{
    double r, i;
    inline void operator +=(const cpx &b){ r+=b.r, i+=b.i;}
    inline cpx operator +(const cpx &b)const{ return (cpx){r+b.r, i+b.i};}
    inline cpx operator -(const cpx &b)const{ return (cpx){r-b.r, i-b.i};}
    inline cpx operator *(const cpx &b)const{ return (cpx){r*b.r-i*b.i, r*b.i+i*b.r};}
    inline cpx operator *(const double b)const{ return (cpx){r*b, i*b};}
    inline cpx operator ~()const{return (cpx){r, -i};}
}a[MAXN<<2], w[MAXN<<2];
ll ans[MAXN << 2];
inline void DFT_(cpx *f, int n)
{
    for(ri i=0, j=0; i<n; ++i)
    {
        if(i>j) swap(f[i], f[j]);
        for(ri k=n>>1; (j^=k)<k; k>>=1);
    }
    for(ri i=1; i<n; i<<=1) for(ri j=0; j<n; j+=i<<1)
            for(ri k=j; k<j+i; ++k)
            {
                cpx t=w[i+k-j]*f[k+i];
                f[k+i]=f[k]-t, f[k]+=t;
            }
}
inline void DFT(cpx *f, int n)
{
    if(n==1) return;
    n>>=1;
    static cpx a[MAXN<<1];
    for(ri i=0; i<n; ++i) a[i]=(cpx){f[i<<1].r, f[i<<1|1].r};
    DFT_(a, n);
    for(ri i=0; i<n; ++i)
    {
        cpx q=~a[(n-i)&(n-1)], x=(a[i]+q)*0.5, y=(a[i]-q)*(cpx){0, -0.5}, t=y*w[n+i];
        f[i]=x+t, f[n+i]=x-t;
    }
}
inline void IDFT(cpx *f, int n)
{
    if(n==1) return;
    reverse(f+1, f+n), n>>=1;
    static cpx a[MAXN<<1];
    for(ri i=0; i<n; ++i)
        a[i]=(f[i]+f[i+n])*0.5 + (f[i]-f[i+n])*(cpx){0, 0.5}*w[n+i];
    DFT_(a, n);
    double k=1.0/n;
    for(ri i=0; i<n; ++i) f[i<<1]=(cpx){a[i].r*k, 0}, f[i<<1|1]=(cpx){a[i].i*k, 0};
}

int n,root,ms,mson[MAXN],sz[MAXN],Size;
bool vis[MAXN];
//root用于标记重心,ms表示树的重心的最大子树的大小,mson[i]记录以i为根最大子树的大小
//sz[i]记录以i为根子树的大小,Size表示当前整棵树的大小,vis[i]表示当前节点是否被分治过
void getroot(int u,int f)//获得重心
{
    sz[u] = 1,mson[u] = 0;
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v] || v == f) continue;//剔除已经被分治过的点
        getroot(v,u);
        sz[u] += sz[v];
        if(sz[v] > mson[u]) mson[u] = sz[v];
    }
    if(Size - sz[u] > mson[u]) mson[u] = Size-sz[u];//把u看作根节点时u的父亲那一部分也算作子树
    if(ms > mson[u]) ms = mson[u],root = u;//更新重心
}
int dis[MAXN],cnt;//dis记录所有节点到重心的距离
int mx;
void getdis(int u,int f,int d)//获得到目标点的距离
{
    dis[++cnt] = d;
    mx = max(mx,d);
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v] || v == f) continue;
        getdis(v,u,d + 1);
    }
}
void cal(int u,int d,int tp)//u表示getdis的起点,d表示u到目标点的距离,tp表示这一次统计出来的答案是合理的还是不合理的
{
    cnt = mx = 0;
    getdis(u,0,d);//算出树中的点到目标点的距离
    int len = 1;
    while(len <= mx*2) len <<= 1;
    for(int i = 0;i < len;++i)
        a[i] = (cpx){0,0};
    for(int i = 1;i <= cnt;++i)
        a[dis[i]].r += 1;
    for(ri i=1; i<len; i<<=1)
    {
        w[i]=(cpx){1, 0};
        for(ri j=1; j<i; ++j)
            w[i+j]=((j&31)==1?(cpx){cos(pi*j/i), sin(pi*j/i)}:w[i+j-1]*w[i+1]);
    }
    DFT(a,len);
    for(ri i = 0;i < len;++i)
        a[i] = a[i] * a[i];
    IDFT(a,len);
    for(int i = 0;i < len;++i)
        ans[i] += (ll)(a[i].r+0.5)*tp;
    for(int i = 1;i <= cnt;++i)
        ans[dis[i]+dis[i]] -= tp;
}
void solve(int u,int ssize)//ssize是当前这棵子树的大小
{
    vis[u] = true;//代码保证每次进来的u都必定是当前这棵树的重心,我们将vis[u]标记为true,表示u点被分治过
    cal(u,0,1);//计算这棵树以u为重心的所有组合,但包括了共用同一条边的情况
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v]) continue;
        cal(v,1,-1);//将共用一条边的不合法情况去除
        ms = INF;//记得每次都要初始化
        Size = sz[v] < sz[u]?sz[v]:(ssize-sz[u]);//因为v实际上可能是u的父亲,故sz需相减
        getroot(v,v);//求出以v为根节点的子树重心
        solve(root,Size);
    }
}
int mlen;
inline void init()
{
    tot = 0,ms = INF,Size = n;
    memset(head,-1,sizeof(int)*(n+1));
    memset(vis,false,sizeof(bool)*(n+1));
    mlen = 1;
    memset(ans,0,sizeof(ll)*(4*n+1));
    while(mlen < 2*n) mlen <<= 1;
}
int main()
{
    euler(MAXN-10);
    while(~scanf("%d",&n))
    {
        init();
        int u,v;
        for(int i = 1;i < n;++i)
        {
            scanf("%d%d",&u,&v);
            addedg(u,v),addedg(v,u);
        }
        getroot(1,1);
        solve(root,Size);
//        IDFT(ans,mlen);
        ll res = 0;
        for(int i = 1;i <= num && prime[i] < mlen;++i)
            res += ans[prime[i]];
        printf("%.6f\n",1.0*res/n/(n-1));
    }
    return 0;
}