POJ1741 Tree 点分治

/*
    题目描述：给出一棵有边权的树，求树里距离和小于等于k的点对的数量。
*/
#pragma warning(disable:4786)
#pragma comment(linker, "/STACK:102400000,102400000")
#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<stack>
#include<queue>
#include<map>
#include<set>
#include<vector>
#include<cmath>
#include<string>
#include<sstream>
#include<bitset>
#define LL long long
#define FOR(i,f_start,f_end) for(int i=f_start;i<=f_end;++i)
#define mem(a,x) memset(a,x,sizeof(a))
#define lson l,m,x<<1
#define rson m+1,r,x<<1|1
using namespace std;
const int INF = 0x3f3f3f3f;
const int mod = 1e9 + 7;
const double PI = acos(-1.0);
const double eps=1e-6;
const int maxn = 1e4 + 5;
int size[maxn] , maxv[maxn] , vis[maxn] , dis[maxn];
int n , k;
int ans , root , num , Max;
int tot;
struct node
{
    int v , w , nt;
}edges[maxn * 2];
int head[maxn];
void init()
{
    tot = 0;
    ans = 0;
    mem(vis , 0);
    mem(head , -1);
}
void add_edge(int u , int v , int w)
{
    edges[tot].v = v;
    edges[tot].w = w;
    edges[tot].nt = head[u];
    head[u] = tot++;
}
void dfssize(int u , int fa)
{
    size[u] = 1;
    maxv[u] = 0;
    for(int i = head[u] ; i != -1 ; i = edges[i].nt){
        int v = edges[i].v;
        if(vis[v] || v == fa)       continue;
        dfssize(v , u);
        size[u] += size[v];
        maxv[u] = max(maxv[u] , size[v]);
    }
}
void dfsroot(int r , int u , int fa)
{
    if(size[r] - size[u] > maxv[u])     maxv[u] = size[r] - size[u];
    if(maxv[u] < Max){
        Max = maxv[u];
        root = u;
    }
    for(int i = head[u] ; i != -1 ; i = edges[i].nt){
        int v = edges[i].v ;
        if(vis[v] || v == fa)       continue;
        dfsroot(r , v , u);
    }
}
void dfsdis(int u , int d , int fa)
{
    dis[num++] = d;
    for(int i = head[u] ; i != -1 ; i = edges[i].nt){
        int v = edges[i].v;
        if(vis[v] || v == fa)       continue;
        dfsdis(v , d + edges[i].w , u);
    }
}
int cal(int u  , int d)     //cal函数巧妙在只要知道子树的根u与u距离root的距离d，并没有真正用到当前的root
{
    int ret = 0;
    num = 0;
    dfsdis(u , d , 0);       //是u,d,0  不是u, 0, 0
    sort(dis , dis + num);
    int i = 0 ,  j = num - 1;
    while(i < j){
        while(dis[i] + dis[j] > k && j > i)     --j;
        ret += (j - i);
        ++i;
    }
    return ret;
}
void dfs(int u)
{
    Max = n;
    dfssize(u , 0);
    dfsroot(u , u , 0);
    ans += cal(root , 0);   //想好哪些地方是u，哪些地方是root！
    vis[root] = 1;
    for(int i = head[root] ; i != -1 ; i = edges[i].nt){
        int v = edges[i].v;
        if(!vis[v]){
            ans -= cal(v , edges[i].w);
            dfs(v);     //虽改变了root，但不影响下一次循环算ans -= cal(v , edges[i].w);
        }
    }
}
int main()
{
    while(scanf("%d %d", &n , &k) != EOF){
        if(!n && !k)        break;
        init();
        int v , u , w;
        for(int i = 1 ; i<n ; i++){
            scanf("%d %d %d", &u , &v , &w);
            add_edge(u , v , w);
            add_edge(v , u , w);
        }
        dfs(1);
        printf("%d\n",ans);
    }
    return 0;
}