【FFT+点分治，树重心】CODECHEF PRIMEDST

通道：https://www.codechef.com/problems/PRIMEDST

题意：一棵树上任取两点，问两点间距离为质数的概率

思路：对于求树上的路径数的题，可以想到要用树分治，我们可以考虑经过一个点的方案数，用fft快速地求出所有方案数，然后减去来自同一颗子树的方案数

代码：

#include <cstdio>
#include <cmath>
#include <cstring>
#include <vector>
#include <algorithm>

using namespace std;

typedef long long ll;

const int MAX_N = 100007; 
const int MAX_M = 100007;
const double PI = acos(-1.0);

struct Complex {
    double r, i;
    Complex(double _r, double _i) {
        r = _r;
        i = _i;
    }
    Complex operator + (const Complex &c) {
        return Complex(c.r + r, c.i + i);
    }
    Complex operator - (const Complex &c) {
        return Complex(r - c.r, i - c.i);
    }
    Complex operator * (const Complex &c) {
        return Complex(c.r * r - c.i * i, c.r * i + c.i * r);
    }
    Complex operator / (const int &c) {
        return Complex(r / c, i / c);
    }
    Complex(){}
};
namespace FFT {
    int rev(int id, int len) {
        int ret = 0;
        for(int i = 0; (1 << i) < len; ++i) {
            ret <<= 1;
            if(id & (1 << i)) ret |= 1;
        }
        return ret;
    }
    Complex A[MAX_M << 3];
    void FFT(Complex *a, int len, int DFT) {
        for(int i = 0; i < len; ++i) A[rev(i, len)] = a[i];
        for(int s = 1; (1 << s) <= len; ++s) {
            int m = (1 << s);
            Complex wm = Complex(cos(PI * DFT * 2 / m), sin(PI * DFT * 2 / m));
            for(int k = 0; k < len; k += m) {
                Complex w = Complex(1, 0);
                for(int j = 0; j < (m >> 1); j++) {
                    Complex t = w * A[k + j + (m >> 1)];
                    Complex u = A[k + j];
                    A[k + j] = u + t;
                    A[k + j + (m >> 1)] = u - t;
                    w = w * wm;
                }
            }
        }
        if(DFT == -1) for(int i = 0; i < len; ++i) A[i] = A[i] / len;
        for(int i = 0; i < len; i++) a[i] = A[i];
    }
};

struct Node {
    int v, nxt;
    Node () {
        
    }
    Node (int _v, int _n) {
        v = _v;
        nxt = _n;
    }
};

int n;
int head[MAX_N], edgecnt;
Node G[MAX_N << 1];
bool p[MAX_N], del[MAX_N];
int prime[MAX_N], primeCnt;

void Clear() {
    edgecnt = 0;
    memset(head, -1, sizeof head);
    memset(del, 0, sizeof del);
}

void init() {
    primeCnt = 0;
    memset(p, 1, sizeof p);
    for (int i = 2; i < MAX_N; ++i) {
        if (p[i]) {
            prime[primeCnt++] = i;
            for (int j = 2 * i; j < MAX_N; j += i)
                p[j] = false;
        }
    }
}

void add(int u, int v){
    G[edgecnt] = Node(v, head[u]);
    head[u] = edgecnt++;
}

int son[MAX_N], opt[MAX_N];
vector<int> alln;

void dfs(int u,int fa) {
    alln.push_back(u);
    son[u] = 1; opt[u] = 0;
    for(int i = head[u]; ~i; i = G[i].nxt) {
        int v = G[i].v;
        if(del[v] || v == fa) continue;
        dfs(v, u);
        son[u] += son[v];
        opt[u] = max(opt[u], son[v]);
    }
}

int getCenter(int u) {
    alln.clear();
    dfs(u, -1);
    int mx = 0, ans = -1;
    int sz = alln.size();
    for(int i = 0; i < sz; ++i)  {
        int v = alln[i];
        if(ans == -1) ans = v, mx = max(opt[v], sz - son[v]);
        else  {
            if(max(opt[v], sz - son[v]) < mx) {
                mx = max(opt[v], sz - son[v]);
                ans = v;
            }
        }
    }
    return ans;
}

int tot, D[MAX_N];
void getDist(int u, int fa, int w) {
    D[tot++] = w;
    for(int i = head[u]; ~i; i = G[i].nxt) {
        int v = G[i].v;
        if(del[v] || v == fa) continue;
        getDist(v, u, w + 1);
    } 
}

int cnt[MAX_N];
ll res[MAX_M << 3];
Complex A[MAX_M << 3];

ll calc() {
    int up = *max_element(D, D + tot);
    memset(cnt, 0, sizeof cnt);
    for (int i = 0; i < tot; ++i) ++cnt[D[i]];
    for (int i = 0; i <= up; ++i) A[i] = Complex(cnt[i], 0);
    int len = 1;
    while (len <= up) len <<= 1; len <<= 1;
    for (int i = up + 1; i < len; ++i) A[i] = Complex(0, 0);
    FFT::FFT(A, len, 1);
    for (int i = 0; i < len; ++i) A[i] = A[i] * A[i];
    FFT::FFT(A, len, -1);
    for (int i = 0; i < len; ++i) res[i] = (ll)(A[i].r + 0.5);
    for (int i = 0; i < tot; ++i) --res[D[i] + D[i]];
    for (int i = 0; i < primeCnt && prime[i] < len; ++i) res[prime[i]] /= 2; 
    ll ans = 0;
    for (int i = 0; i < primeCnt && prime[i] < len; ++i) ans += res[prime[i]];
    return ans;
}

ll ans;
void solve(int u) {
    u = getCenter(u);
    tot = 0; getDist(u, -1, 0);
    ans += calc();
    for(int i = head[u]; ~i; i = G[i].nxt) {
        int v = G[i].v;
        if(del[v]) continue;
        tot = 0;
        getDist(v, u, 1);
        ans -= calc();
    }
    del[u] = true;
    for(int i = head[u]; ~i; i =G[i].nxt) {
        int v = G[i].v;
        if(del[v]) continue;
        solve(v);
    }
}

int main() {
    init();
    while (1 == scanf("%d", &n)) {
        Clear();
        for (int i = 1; i < n; ++i) {
            int u, v;
            scanf("%d%d", &u, &v);
            add(u, v), add(v, u);
        }
        ans = 0;
        solve(1);
        double an = 1. * ans * 2 / n / (n - 1);
        printf("%.8f\n", an);
    }
    return 0;
}