【树分】POJ - 1741 - 树中点对统计

【论文链接】分治算法在树的路径问题中的应用

---

POJ - 1741 - 树中点对统计

题目链接<http://poj.org/problem?id=1741>

---

#include<iostream>
#include<stdio.h>
#include<algorithm>
#include<string.h>
#define ll long long
using namespace std;
const int N=1e4+7;
struct Edge{
    int v,w,nxt;
    Edge(int v=0,int w=0,int nxt=0):v(v),w(w),nxt(nxt){}
}e[N*2];
int sz[N],maxn[N],dis[N];
bool vis[N];
int n,k,rt,Max,tot,ans;
int p[N],edn;
void add(int u,int v,int w){
    e[++edn]=Edge(v,w,p[u]);p[u]=edn;
    e[++edn]=Edge(u,w,p[v]);p[v]=edn;
}
void dfssz(int u,int f){
    sz[u]=1;maxn[u]=0;
    for(int i=p[u];~i;i=e[i].nxt){
        int v=e[i].v;
        if(v==f||vis[v]) continue;
        dfssz(v,u);
        sz[u]+=sz[v];
        maxn[u]=max(maxn[u],sz[v]);
    }
}
void dfsrt(int u,int f,int r){
    maxn[u]=max(maxn[u],sz[r]-sz[u]);
    if(maxn[u]<Max) Max=maxn[u],rt=u;
    for(int i=p[u];~i;i=e[i].nxt){
        int v=e[i].v;
        if(v==f||vis[v]) continue;
        dfsrt(v,u,r);
    }
}
void dfsdis(int u,int f,int d){
    dis[++tot]=d;
    for(int i=p[u];~i;i=e[i].nxt){
        int v=e[i].v;
        if(v==f||vis[v]) continue;
        dfsdis(v,u,d+e[i].w);
    }
}
int cal(int u,int d){
    int res=0;
    tot=0;
    dfsdis(u,u,d);
    sort(dis+1,dis+1+tot);
    int i=1,j=tot;
    while(i<j){
        while(j>i&&dis[i]+dis[j]>k) j--;
        res+=j-i;
        i++;
    }
    return res;
}
void dfs(int u){
    Max=n;
    dfssz(u,0);
    dfsrt(u,0,u);
    ans+=cal(rt,0);
    vis[rt]=true;
    for(int i=p[rt];~i;i=e[i].nxt){
        int v=e[i].v;
        if(vis[v]) continue;
        ans-=cal(v,e[i].w);
        dfs(v);
    }
}
int main()
{
    while(scanf("%d%d",&n,&k)!=EOF){
        if(n==0&&k==0) break;
        memset(p,-1,sizeof(p));edn=-1;ans=0;
        memset(vis,false,sizeof(vis));
        int u,v,w;
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
        }
        dfs(1);
        printf("%d\n",ans);
    }
}