CodeChef - PRIMEDST (点分治/DSU+FFT)

CodeChef - PRIMEDST (点分治/DSU+FFT)

题目的本质是要求每种距离的点对的个数

考虑使用点分治+FFT来维护

对于当前点分根以及周围的所有点,构造\(f(x)=\sum x^{dep_x}\)求平方即可，对于每颗子树容斥

复杂度就是所有的\(Size\)之和乘上\(\log n\)，即\(n\log^2 n\)

最后把距离是质数的算出来即可

(代码太丑了)

#include<bits/stdc++.h>
using namespace std;

//#define double long double

#define reg register
typedef long long ll;
#define rep(i,a,b) for(reg int i=a,i##end=b;i<=i##end;++i)
#define drep(i,a,b) for(reg int i=a,i##end=b;i>=i##end;--i)

template <class T> inline void cmin(T &a,T b){ ((a>b)&&(a=b)); } 
template <class T> inline void cmax(T &a,T b){ ((a<b)&&(a=b)); } 

char IO;
int rd(){
    int s=0,f=0;
    while(!isdigit(IO=getchar())) if(IO=='-') f=1;
    do s=(s<<1)+(s<<3)+(IO^'0');
    while(isdigit(IO=getchar()));
    return f?-s:s;
}

const double PI=acos(-1);

const int N=(1<<18)+10,P=1e6+3;
const int g=3;

bool be;
int n;
struct Edge{
    int to,nxt;
} e[N<<1];
int head[N],ecnt;
void AddEdge(int u,int v) {
    e[++ecnt]=(Edge){v,head[u]};
    head[u]=ecnt;
}
#define erep(u,i) for(int i=head[u];i;i=e[i].nxt)

int rt,mi,sz[N],vis[N];
void GetSize(int u,int f) {
    sz[u]=1;
    erep(u,i) {
        int v=e[i].to;
        if(v==f||vis[v]) continue;
        GetSize(v,u);
        sz[u]+=sz[v];
    }
}
void FindRt(int n,int u,int f) {
    int ma=n-sz[u];
    erep(u,i) {
        int v=e[i].to;
        if(v==f||vis[v]) continue;
        ma=max(ma,sz[v]);
        FindRt(n,v,u);
    }
    if(ma<mi) mi=ma,rt=u;
}

ll ans[N],sum;
int notpri[N];
int rev[N];
struct Cp{
    double x,y;
    Cp(){}
    Cp(double _x,double _y){ x=_x,y=_y; }
    Cp operator + (const Cp t){ return Cp(x+t.x,y+t.y); }
    Cp operator - (const Cp t){ return Cp(x-t.x,y-t.y); }
    Cp operator * (const Cp t){ return Cp(x*t.x-y*t.y,x*t.y+y*t.x); }
}a[N],b[N],c[N];

struct Node{
    int x,maxd;
    bool operator < (const Node __) const {
        return maxd<__.maxd;
    }
}son[N];

void FFT(int n,Cp *a,int f){
    rep(i,1,n-1) if(i<rev[i]) swap(a[i],a[rev[i]]);
    for(reg int i=1;i<n;i<<=1) {
        Cp w(cos(PI/i),f*sin(PI/i));
        for(reg int l=0;l<n;l+=i*2) {
            Cp e(1,0);
            for(reg int j=l;j<l+i;j++,e=e*w) {
                Cp t=a[j+i]*e;
                a[j+i]=a[j]-t;
                a[j]=a[j]+t;
            }
        }
    }
    if(f==-1) rep(i,0,n-1) a[i].x/=n;
}

int maxd,cnt[N],nc[N];
void dfs(int u,int f,int d,int k) {
    cmax(maxd,d);
    cnt[d]+=k;
    sz[u]=1;
    erep(u,i) {
        int v=e[i].to;
        if(v==f||vis[v]) continue;
        dfs(v,u,d+1,k);
        sz[u]+=sz[v];
    }
}

void Solve(int u){ 
    //cout<<"Begin Solve #"<<u<<endl;
    vis[u]=1;
    int sonc=0;
    maxd=0;
    erep(u,i) {
        int v=e[i].to;
        if(vis[v]) continue;
        maxd=0;
        dfs(v,u,1,0);
        son[++sonc]=(Node){v,maxd};
    }
    sort(son+1,son+sonc+1);
    nc[0]++;
    rep(t,1,sonc) {
        dfs(son[t].x,u,1,1);
        //cout<<"son="<<son[t].x<<endl;
        int R=1,c=-1;
        while(R<=maxd*2) R<<=1,c++;
        rep(i,1,R) rev[i]=(rev[i>>1]>>1)|((i&1)<<c);
        rep(i,0,R) a[i].x=cnt[i],b[i].x=nc[i];
        FFT(R,a,1),FFT(R,b,1);
        rep(i,0,R) a[i]=a[i]*b[i];
        FFT(R,a,-1);
        rep(i,0,R) ans[i]+=ll(a[i].x+0.5);
        rep(i,0,maxd) nc[i]+=cnt[i],cnt[i]=0;
        rep(i,0,R) a[i].x=a[i].y=b[i].x=b[i].y=0;
    }
    rep(i,0,maxd) nc[i]=0;
    erep(u,i) {
        int v=e[i].to;
        if(vis[v]) continue;
        mi=1e9,FindRt(sz[v],v,u);
        Solve(rt);
    }
}

bool ed;

int main(){
    n=rd();
    rep(i,2,n) {
        int u=rd(),v=rd();
        AddEdge(u,v),AddEdge(v,u);
    }
    GetSize(1,0),mi=1e9,FindRt(n,1,0);
    Solve(rt);
    rep(i,2,n) if(!notpri[i]) {
        sum+=ans[i];
        for(reg int j=i+i;j<=n;j+=i) notpri[j]=1;
    }
    //rep(i,1,n) cout<<ans[i]<<' '; puts("");
    printf("%.7lf\n",1.0*sum/(1.0*n*(n-1)/2));
}