POJ 1741 [点分治][树上路径问题]

/*
不要低头，不要放弃，不要气馁，不要慌张
题意：
给一棵有n个节点的树，每条边都有一个正权值，求一共有多少个点对使得它们之间路的权值和小于给定的k.
思路：
《分治算法在树的路径问题中的应用》
*/

#include<stdio.h>
#include<vector>
#include<string.h>
#include<algorithm>
#define N 10050
using namespace std;
struct edge{
    int id;
    long long w;
    bool im;
    edge *next;
};
struct st{
    st(){}
    st(int a,long long b,int c){
        id=a;dis=b;iid=c;
    }
    int id,iid;
    long long dis;
};
vector<st>mv;
edge edges[N*];
edge *adj[N];
long long ans,kk;
int ednum;
inline void addedge(int a,int b,long long c){
    edge *tmp=&edges[ednum++];
    tmp->id=b;
    tmp->w=c;
    tmp->im=;
    tmp->next=adj[a];
    adj[a]=tmp;
}
int zong,next_val,next;
bool vis[N];
int siz[N],fa[N];
long long dis[N];
void dfs(int pos,int dep){
    vis[pos]=;
    siz[pos]=;
    for(edge *it=adj[pos];it;it=it->next){
        if(it->im&&!vis[it->id]){
            dfs(it->id,dep+);
            siz[pos]+=siz[it->id];
        }
    }
}
void dfs2(int pos,int dep){
    int my_next=-;
    vis[pos]=;
    for(edge *it=adj[pos];it;it=it->next){
        if(it->im&&!vis[it->id]){
            my_next=max(my_next,siz[it->id]);
        }
    }
    my_next=max(my_next,zong-siz[pos]);
    if(next_val>my_next){
        next=pos;
        next_val=my_next;
    }
    for(edge *it=adj[pos];it;it=it->next){
        if(it->im&&!vis[it->id]){
            dfs2(it->id,dep+);
        }
    }
}
bool cmp1(st a,st b){
    if(a.dis!=b.dis)return a.dis<b.dis;
    else return a.iid<b.iid;
}
bool cmp2(st a,st b){
    if(a.id!=b.id)return a.id<b.id;
    else if(a.dis!=b.dis)return a.dis<b.dis;
    else return a.iid<b.iid;
}
inline void del(int a,int b){
    for(edge *it=adj[a];it;it=it->next){
        if(it->id==b){
            it->im=;
            return;
        }
    }
}
void dfs3(int pos,int dep){
    vis[pos]=;
    if(!dep)dis[pos]=;
    for(edge *it=adj[pos];it;it=it->next){
        if(it->im&&!vis[it->id]){
            if(!dep)fa[it->id]=it->id;
            else fa[it->id]=fa[pos];
            dis[it->id]=dis[pos]+it->w;
            mv.push_back(st(fa[it->id],dis[it->id],it->id));
            dfs3(it->id,dep+);
        }
    }
}
void solve(int pos){
    mv.clear();
    memset(vis,,sizeof(vis));
    dfs(pos,);
    zong=siz[pos];
    if(zong<=)return;
    memset(vis,,sizeof(vis));
    next_val=;
    dfs2(pos,);
    memset(vis,,sizeof(vis));
    dfs3(next,);
    int n=mv.size();
    sort(mv.begin(),mv.end(),cmp1);
    for(int i=;i<n;i++){
        if(mv[i].dis>kk)break;
        ans++;
        int l=i+,r=n-;
        while(l<=r){
            int mid=(l+r)>>;
            if(mv[i].dis+mv[mid].dis<=kk)l=mid+;
            else r=mid-;
        }
        ans+=r-i;
    }
    sort(mv.begin(),mv.end(),cmp2);
    int st=;
    for(int i=;i<n;i++){
        if(mv[i].id!=mv[st].id){
            for(int j=st;j<i;j++){
                int l=j+,r=i-;
                while(l<=r){
                    int mid=(l+r)>>;
                    if(mv[j].dis+mv[mid].dis<=kk)l=mid+;
                    else r=mid-;
                }
                ans-=r-j;
            }
            st=i;
        }
    }
    for(int j=st;j<n;j++){
        int l=j+,r=n-;
        while(l<=r){
            int mid=(l+r)>>;
            if(mv[j].dis+mv[mid].dis<=kk)l=mid+;
            else r=mid-;
        }
        ans-=r-j;
    }
    vector<int>mmv;
    for(edge *it=adj[next];it;it=it->next){
        if(it->im&&siz[it->id]>){
            mmv.push_back(it->id);
            it->im=;
            del(it->id,next);
        }
    }
    for(int i=;i<mmv.size();i++)solve(mmv[i]);
}
int main()
{
    int n;
    while(scanf("%d%lld",&n,&kk)!=EOF){
        ans=;
        if(!n)break;
        memset(adj,NULL,sizeof(adj));
        ednum=;
        for(int i=;i<n;i++){
            int a,b;
            long long c;
            scanf("%d%d%lld",&a,&b,&c);
            addedge(a,b,c);
            addedge(b,a,c);
        }
        solve();
        printf("%lld\n",ans);
    }
}