CC PRIMEDST(点分治+FFT)

• 题意

给你一颗树，求距离为质数的点对个数。

• 题解

裸题。。。
顺便打一份点分治的板，以前打错板从来没错过。。。

#include<iostream>
#include<cstdio>
#include<string>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<vector>
using namespace std;
const int Maxn=8e4+50;
const double PI2=2.0*acos(-1.0);

inline int read()
{
    char ch=getchar();int i=0,f=1;
    while(!isdigit(ch)){if(ch=='-')f=-1;ch=getchar();}
    while(isdigit(ch)){i=(i<<1)+(i<<3)+ch-'0';ch=getchar();}
    return i*f;
}

int Prime[Maxn],tot,pr[Maxn];
int ans,n,rt,maxdis,total,sze[Maxn],mx[Maxn],last[Maxn],before[Maxn*2],to[Maxn*2],ecnt=1,vis[Maxn],cnt1[Maxn],cnt2[Maxn],C,dis[Maxn];

inline void sieve(int lim)
{
    for(int i=2;i<=lim;i++)
    {
        if(!pr[i]){Prime[++tot]=i;pr[i]=i;}
        for(int j=1;j<=tot;j++)
        {
            int k=i*Prime[j];
            if(k>lim)break;
            pr[k]=Prime[j];
            if(!(i%Prime[j]))break;
        }
    }
}

struct Complex
{
    double r,i;
    Complex(double r=0,double i=0):r(r),i(i){}
    friend inline Complex operator +(const Complex &a,const Complex &b){return Complex(a.r+b.r,a.i+b.i);}
    friend inline Complex operator -(const Complex &a,const Complex &b){return Complex(a.r-b.r,a.i-b.i);}
    friend inline Complex operator *(const Complex &a,const Complex &b){return Complex(a.r*b.r-a.i*b.i,a.r*b.i+a.i*b.r);}
}f1[Maxn],f2[Maxn];

struct FFT
{
    int pos[Maxn],k;
    inline void init(int len)
    {
        for(k=1;k<=len;k<<=1);
        for(int i=1;i<k;i++)pos[i]=(i&1)?((pos[i>>1]>>1)^(k>>1)):(pos[i>>1]>>1);
    }
    inline void dft(Complex *a,int o)
    {
        for(int i=1;i<k;i++)if(pos[i]>i)swap(a[pos[i]],a[i]);
        for(int m1=1;m1<k;m1<<=1)
        {
            int m2=m1<<1;
            double t=PI2/(double)m2*o;
            Complex wn(cos(t),sin(t));
            for(int i=0;i<k;i+=m2)
            {
                Complex w(1,0);
                for(int j=0;j<m1;j++)
                {
                    Complex &A=a[i+j],&B=a[i+j+m1],t=w*B;
                    B=(A-t);
                    A=(A+t);
                    w=w*wn;
                }
            }
        }
    }
    inline void mui(int *a,int *b,Complex *c,Complex *d,Complex *e)
    {
        for(int i=0;i<k;i++)c[i].r=a[i],d[i].r=b[i],c[i].i=d[i].i=0;
        dft(c,1);dft(d,1);
        for(int i=0;i<k;i++)e[i]=c[i]*d[i];
        dft(e,-1);
        for(int i=0;i<k;i++)e[i].r/=k;
    }
}fft;

inline void add(int x,int y){before[++ecnt]=last[x];last[x]=ecnt;to[ecnt]=y;}

inline void getdep(int now,int val,int f)
{
    maxdis=max(maxdis,val);
    if(val)cnt2[val]++;
    sze[now]=1;
    for(int e=last[now];e;e=before[e])
    {
        int v=to[e];
        if(v^f&&!vis[v])
        {
            getdep(v,val+1,now);
            sze[now]+=sze[v];
        }
    }
}

inline void dfs(int now,int f)
{
    sze[now]=1,mx[now]=0;
    for(int e=last[now];e;e=before[e])
    {
        int v=to[e];
        if(v^f&&!vis[v])
        {
            dfs(v,now);
            sze[now]+=sze[v];
            mx[now]=max(mx[now],sze[v]);
        }
    }
    mx[now]=max(mx[now],total-sze[now]);
    if(mx[now]<mx[rt])rt=now;
}

inline void del(int now,int val,int f=0)
{
    if(val)cnt1[val]--;
    for(int e=last[now];e;e=before[e])
    {
        int v=to[e];
        if(v^f&&!vis[v])del(v,val+1,now);
    }
}

inline void solve(int now)
{
    vis[now]=1;maxdis=0;cnt1[0]=1;
    for(int e=last[now];e;e=before[e])
    {
        int v=to[e];
        if(!vis[v])
        {
            getdep(v,1,now);
            fft.init(2*maxdis+1);
            fft.mui(cnt1,cnt2,f1,f2,f2);
            for(int i=1;i<=tot&&Prime[i]<=2*maxdis;i++){ans+=(int)(f2[Prime[i]].r+0.5);}
            for(int i=1;i<=maxdis;i++)cnt1[i]+=cnt2[i],cnt2[i]=0;
        }
    }
    del(now,0);
    for(int e=last[now];e;e=before[e])
    {
        int v=to[e];
        if(!vis[v])
        {
            total=mx[0]=sze[v];rt=0;
            dfs(v,now);solve(rt);
        }
    }
}

int main()
{
    n=read();
    sieve(n);
    for(int i=1;i<n;i++){int x=read(),y=read();add(x,y);add(y,x);}
    total=mx[0]=n;rt=0;
    dfs(1,0);
    solve(rt);
    printf("%.8f",(double)ans/((double)n*(double)(n-1)/2.0));
}