【BZOJ 1468】Tree （点分治）

Problem

给你一棵TREE,以及这棵树上边的距离。
问有多少对点它们两者间的距离小于等于K

Solution

点分第一道【捂脸】
点分的基本操作就是找重心，做与这个重心相关的信息

容斥…

Code

#include <cstdio>
#include <cstring>
#include <queue>
#include <algorithm>
using namespace std;
#define N 40010
#define inf 0x3f3f3f3f
int n,m,num=0,cnt=0,sum,root,k,ans=0,d[N],sz[N],mx[N],vis[N],h[N],q[N];
struct node{int to,z,next;}mp[N<<1];
inline char gc(){
    static char buf[1<<16],*S,*T;
    if(S==T){T=(S=buf)+fread(buf,1,1<<16,stdin);if(T==S) return EOF;}
    return *S++;
}
inline int read(){
    int x=0,f=1;char ch=gc();
    while(ch<'0' || ch>'9'){if(ch=='-') f=-1;ch=gc();}
    while('0'<=ch && ch<='9') x=x*10+ch-'0',ch=gc();
    return x*f;
}
inline void insert(int x,int y,int z){
    mp[++num].to=y;mp[num].z=z;mp[num].next=h[x];h[x]=num;
    mp[++num].to=x;mp[num].z=z;mp[num].next=h[y];h[y]=num;
}
void getroot(int x,int fa){
    sz[x]=1;mx[x]=0;
    for(int i=h[x];i;i=mp[i].next){
        int y=mp[i].to;if(y==fa || vis[y]) continue;
        getroot(y,x);sz[x]+=sz[y];mx[x]=max(mx[x],sz[y]);
    }
    mx[x]=max(mx[x],sum-sz[x]);
    if(mx[x]<mx[root]) root=x;
}
void dist(int x,int fa){
    q[++cnt]=d[x];
    for(int i=h[x];i;i=mp[i].next){
        int y=mp[i].to;if(y==fa || vis[y]) continue;
        d[y]=d[x]+mp[i].z;dist(y,x);
    }
}
int calc(int rt,int val){
    d[rt]=val;cnt=0;dist(rt,0);
    sort(q+1,q+cnt+1);
    int l=1,r=cnt,res=0;
    while(l<r){
        if(q[l]+q[r]>k) r--;
        else res+=r-l,l++;
    }return res;
}
void solve(int x){
    ans+=calc(x,0);vis[x]=1;
    for(int i=h[x];i;i=mp[i].next){
        int y=mp[i].to;if(vis[y]) continue;
        ans-=calc(y,mp[i].z);
        sum=sz[y];root=0;getroot(y,0);solve(y);
    }
}
int main(){
    n=read();
    for(int i=1;i<n;i++){
        int x=read(),y=read(),z=read();
        insert(x,y,z);
    }k=read();
    mx[0]=inf;sum=n;getroot(1,0);
    sum=n;solve(root);
    printf("%d\n",ans);
    return 0;
}