CodeChef Prime Distance On Tree（点分治+FFT）

题目链接

题目描述：统计树上路径中有多少长度为素数的路径

分析：
树上路径统计问题，考虑点分治
没有什么特别的性质，所以我们只能统计所有路径的长度
不过朴素算法统计路径需要$n^2$的复杂度，点分治还不如LCA好用nai。。。

我们在点分治的时候，当前重心为$rt$，当前处理子树为$x$
我们dfs一遍子树x，那么子树x的所有路径就可以和已经处理过的子树中的所有路径依次组成新路径（经过$rt$）
相当于从两个集合中各选一个数组合出素数的方案数（组合问题？）
总而言之，就是一个类似多项式乘法的东西，用FFT解决

FFT，多项式乘法和生成函数的关系

因为要进行多次FFT，而主n次单位根的计算公式：

$$e^{2\pi i \over n}=(cos({2\pi \over n}),sin({2\pi \over n})i)$$

中有一个$/n$，不同的式子n不同，所以预处理不大现实了，所以我们直接在FFT中处理$\omega$即可

node wn(cos(2.0*opt*Pi/i),sin(2.0*opt*Pi/i));

tip

前辈表示：NTT TLE，FFT跑的超级快……

计数会爆int哦

一开始T了，深搜好像是比广搜费时，改掉
真费劲。。。

#include<cstdio>
#include<cstring>
#include<iostream>
#include<cmath>
#define ll long long

using namespace std;

const int N=100005;
const double Pi=acos(-1.0);
struct node{
    double x,y;
    node (double xx=0,double yy=0) {
        x=xx;y=yy;
    }
}a[140000]; 

node operator +(const node &A,const node &B) {return node(A.x+B.x,A.y+B.y);}
node operator -(const node &A,const node &B) {return node(A.x-B.x,A.y-B.y);}
node operator *(const node &A,const node &B) {return node(A.x*B.x-A.y*B.y,A.x*B.y+A.y*B.x);}

struct E{
    int y,nxt;
}way[N<<1];
int st[N],tot=0,n,root,size[N],f[N];
int deep[N],pre[N],wei,tou,d[N];
bool prime[N],vis[N];
ll ans[N];

void add(int u,int w) {
    tot++;way[tot].y=w;way[tot].nxt=st[u];st[u]=tot;
    tot++;way[tot].y=u;way[tot].nxt=st[w];st[w]=tot;
}

void prepare() {       //E氏筛法 
    memset(prime,1,sizeof(prime)); prime[1]=0;
    for (int i=2;i<=n;i++) if (prime[i])
        for (int j=i+i;j<=n;j+=i) prime[j]=0; 
}

void findroot(int now) {
    tou=wei=0;
    d[++wei]=now;
    pre[now]=0;
    while (tou<wei) {
        int now=d[++tou];
        size[now]=1; f[now]=0;
        for (int i=st[now];i;i=way[i].nxt)
            if (way[i].y!=pre[now]&&!vis[way[i].y]) {
                pre[way[i].y]=now;
                d[++wei]=way[i].y;
            }
    }
    root=0;
    for (int i=wei;i>=1;i--) {
        int p=pre[d[i]],u=d[i];
        size[p]+=size[u];
        f[p]=max(f[p],size[u]);
        f[u]=max(f[u],wei-size[u]);
        if (f[u]<f[root]) root=u;
    }
} 

void FFT(int n,node *a,int opt) {
    int i,j=0,k;
    for (i=0;i<n;i++) {
        if (i>j) swap(a[i],a[j]);
        for (int l=n>>1;(j^=l)<l;l>>=1);
    }
    for (int i=2;i<=n;i<<=1) {
        int m=i>>1;
        node wn(cos(2.0*opt*Pi/i),sin(2.0*opt*Pi/i));
        for (j=0;j<n;j+=i) {
            node w(1,0);
            for (k=0;k<m;k++,w=w*wn) {
                node z=a[j+m+k]*w;
                a[j+m+k]=a[j+k]-z;
                a[j+k]=a[j+k]+z;
            }
        }
    }
} 

void bfs(int now) {           //bfs 
    tou=wei=0;
    d[++wei]=now;
    pre[now]=0;
    while (tou<wei) {
        int now=d[++tou];
        for (int i=st[now];i;i=way[i].nxt)
            if (way[i].y!=pre[now]&&!vis[way[i].y]) {
                deep[way[i].y]=deep[now]+1;
                pre[way[i].y]=now;
                d[++wei]=way[i].y;
            }
    }
} 

void cal(int now,int dep,int opt) {
    deep[now]=dep;
    bfs(now);
    int fn=1;
    while (fn<=deep[d[wei]]*2) fn<<=1;
    for (int i=0;i<=fn;i++) a[i]=node(0.0,0.0);
    for (int i=1;i<=wei;i++) a[deep[d[i]]].x++;
    FFT(fn,a,1);
    for (int i=0;i<=fn;i++) a[i]=a[i]*a[i];
    FFT(fn,a,-1);
    for (int i=0;i<=fn;i++) {
        ll t=(ll)(a[i].x/fn+0.5);
        ans[i]+=(ll)t*opt;
    }
} 

void solve(int now) {
    cal(now,0,1);
    vis[now]=1;
    for (int i=st[now];i;i=way[i].nxt)
        if (!vis[way[i].y]) {
            cal(way[i].y,1,-1);
            findroot(way[i].y);
            solve(root);
        }
}

int main()
{
    scanf("%d",&n);
    prepare();
    for (int i=1;i<n;i++) {
        int u,w;
        scanf("%d%d",&u,&w);
        add(u,w);
    }

    f[0]=N; findroot(1);
    solve(root);

    ll X=0;
    ll Y=(ll)n*(n-1);                               //点分治计算出来的路径端点是排列 
    for (int i=2;i<=n;i++) if (prime[i]) X+=ans[i];
    printf("%0.6lf",(double)X/(double)Y);
    return 0;
}