CodeChef - PRIMEDST Prime Distance On Tree

题意：

给定一个有 $n$ 个结点的树，边权为 $1$，随机选取一对结点，问结点间距离为质数的概率。$(n\le 5\times 10^4)$

链接：

https://vjudge.net/problem/CodeChef-PRIMEDST

解题思路：

路径问题，用点分。$cal$ 函数的求解，$fft$ 处理出所有路径长度，再判断即可。

参考代码：

#include<bits/stdc++.h>

using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
#define sz(a) ((int)a.size())
#define pb push_back
#define lson (rt << 1)
#define rson (rt << 1 | 1)
#define gmid (l + r >> 1)
const int maxn = 5e5 + 5;
const int inf = 0x3f3f3f3f;
const int mod = 1e9 + 7;

struct FFT{

    struct cp{

        double x, y;
        cp operator + (cp o) { return {x + o.x, y + o.y};}
        cp operator - (cp o) { return {x - o.x, y - o.y};}
        cp operator * (cp o) { return {x * o.x - y * o.y, x * o.y + y * o.x};}
    } w[maxn], a[maxn], b[maxn];
    const double PI = acos(-1.0);
    int rev[maxn], len, n, m;
    void init(ll aa[], int nn, ll bb[], int mm){

        n = nn, m = mm;
        for(len = 1; len <= n + m; len <<= 1);
        for(int i = 0; i < len; ++i) a[i] = b[i] = {0, 0};
        for(int i = 0; i <= n; ++i) a[i].x = aa[i];
        for(int i = 0; i <= m; ++i) b[i].x = bb[i];
        for(int i = 0; i < len / 2; ++i){

            double sita = 2 * PI * i / len;
            double c = cos(sita), s = sin(sita);
            w[i] = {c, s}, w[i + len / 2] = {-c, -s};
        }
        for(int i = 0; i < len; ++i) rev[i] = (i & 1) * (len >> 1) + (rev[i >> 1] >> 1);
    }
    void fft(cp x[], int len, int on){

        for(int i = 0; i < len; ++i) if(i < rev[i]) swap(x[i], x[rev[i]]);
        for(int i = 2; i <= len; i <<= 1){

            int wn = len / i, d = i >> 1;
            for(int j = 0; j < len; j += i){

                for(int k = j, wk = 0; k < j + d; ++k, wk += wn){

                    cp a = x[k], b = w[wk] * x[k + d];
                    x[k] = a + b, x[k + d] = a - b;
                }
            }
        }
        if(on == -1){

            reverse(x + 1, x + len);
            for(int i = 0; i < len; ++i) x[i].x /= len;
        }
    }
    void work(ll c[]){

        fft(a, len, 1), fft(b, len, 1);
        for(int i = 0; i < len; ++i) a[i] = a[i] * b[i];
        fft(a, len, -1);
        for(int i = 0; i <= n + m; ++i) c[i] = (ll)(a[i].x + 0.5);
    }
} fft;
vector<int> G[maxn];
int tag[maxn], siz[maxn], vis[maxn], dis[maxn];
int n, tot, tn, rmn, rt;
ll a[maxn], ans;

void getRt(int u, int f){

    int mx = 0; siz[u] = 1;
    for(auto v : G[u]){

        if(v == f || vis[v]) continue;
        getRt(v, u);
        siz[u] += siz[v];
        mx = max(mx, siz[v]);
    }
    mx = max(mx, tn - siz[u]);
    if(mx < rmn) rmn = mx, rt = u;
}

void dfs(int u, int f, int d){

    dis[++tot] = d;
    for(auto v : G[u]){

        if(v == f || vis[v]) continue;
        dfs(v, u, d + 1);
    }
}

void cal(int u, int d, int flg){

    tot = 0;
    dfs(u, 0, d);
    int mx = *max_element(dis + 1, dis + 1 + tot);
    for(int i = 0; i <= mx; ++i) a[i] = 0;
    // cout << u << " " << d << " " << flg << " " << tot << " " << mx << " ?? " << endl;
    for(int i = 1; i <= tot; ++i) ++a[dis[i]];
    // for(int i = 0; i <= mx; ++i) cout << i << " xx " << a[i] <<endl;
    fft.init(a, mx, a, mx);
    fft.work(a);
    // for(int i = 0; i <= mx + mx; ++i) cout << i << " yy " << a[i] << endl;
    for(int i = 0; i <= mx + mx; ++i) ans += flg * a[i] * !tag[i];
    // cout << ans << " xxxxx " << endl;
    // cout << "------" << endl;
}

void dfz(int u){

    vis[u] = 1;
    cal(u, 0, 1);
    for(auto v : G[u]){

        if(vis[v]) continue;
        cal(v, 1, -1);
        tn = siz[v], rmn = inf, getRt(v, u);
        dfz(rt);
    }
    vis[u] = 0;
}

int main(){

    ios::sync_with_stdio(0); cin.tie(0);
    cin >> n;
    tag[0] = tag[1] = 1;
    for(int i = 2; i <= n + n; ++i){

        if(!tag[i]){

            for(int j = i + i; j <= n + n; j += i) tag[j] = 1;
        }
    }
    for(int i = 1; i < n; ++i){

        int u, v; cin >> u >> v;
        G[u].pb(v), G[v].pb(u);
    }
    tn = n, rmn = inf, getRt(1, 0);
    dfz(rt);
    ll tot = n * 1ll * (n - 1);
    // cout << tot << " " << ans << endl;
    cout << ((double)ans / tot) << endl;
    return 0;
}