2018.07.20 洛谷P4178 Tree(点分治)

时间:2024-01-14 16:08:50

传送门

又一道点分治。

直接维护子树内到根的所有路径长度,然后排序+双指针统计答案。

代码如下:

#include<bits/stdc++.h>
#define N 40005
using namespace std;
inline int read(){
    int ans=0;
    char ch=getchar();
    while(!isdigit(ch))ch=getchar();
    while(isdigit(ch))ans=(ans<<3)+(ans<<1)+(ch^48),ch=getchar();
    return ans;
}
int rt,sum,ans,n,k,tot=0,cnt=0,msiz[N],siz[N],d[N],first[N];
bool vis[N];
struct Node{int v,next,w;}e[N<<1];
inline void add(int u,int v,int w){e[++tot].v=v,e[tot].w=w,e[tot].next=first[u],first[u]=tot;}
inline void getroot(int p,int fa){
    siz[p]=1,msiz[p]=0;
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(v==fa||vis[v])continue;
        getroot(v,p),siz[p]+=siz[v];
        if(msiz[p]<siz[v])msiz[p]=siz[v];
    }
    if(sum-siz[p]>msiz[p])msiz[p]=sum-siz[p];
    if(msiz[p]<msiz[rt])rt=p;
}
inline void getdis(int p,int fa,int dis){
    d[++cnt]=dis;
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(vis[v]||v==fa)continue;
        getdis(v,p,dis+e[i].w);
    }
}
inline int calc(int p,int v){
    int ret=0;
    cnt=0;
    getdis(p,0,v);
    sort(d+1,d+cnt+1);
    int l=1,r=cnt;
    while(l<r){
        while(d[l]+d[r]>k&&r>l)--r;
        if(d[l]+d[r]<=k)ret+=r-l;
        ++l;
    }
    return ret;
}
inline void solve(int p){
    ans+=calc(p,0);
    vis[p]=true;
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(vis[v])continue;
        ans-=calc(v,e[i].w);
        sum=siz[v],rt=0,getroot(v,0);
        solve(rt);
    }
}
int main(){
    n=read();
    for(int i=1;i<n;++i){
        int u=read(),v=read(),w=read();
        add(u,v,w),add(v,u,w);
    }
    k=read();
    msiz[rt=0]=sum=n,getroot(1,0),solve(rt);
    printf("%d",ans);
    return 0;
}