POJ - 1741 Tree 树的分治

题目大意：给出一棵N个点的树，每条边都有相应的权值。
先给出K，要求你找出权值小于等于k的（u,v）对

解题思路：具体的思路可以参考漆子超的《分治算法在树的路径问题中的应用》这篇论文。

#include<cstdio>
#include<algorithm>
#include<vector>
using namespace std;
#define maxn 10010
vector<int> Node[maxn], W[maxn], dep;
int Sum[maxn], dp[maxn], vis[maxn], d[maxn];
int n, k, root, size, Ans;

void dfs(int cur, int fa) {
    Sum[cur] = 1;
    dp[cur] = 0;
    for(int i = 0; i < Node[cur].size(); i++) {
        if(!vis[Node[cur][i]] && Node[cur][i] != fa) {
            dfs(Node[cur][i], cur);
            Sum[cur] += Sum[Node[cur][i]];
            dp[cur] = max(dp[cur], Sum[Node[cur][i]]);
        }
    }
    dp[cur] = max(dp[cur], size - Sum[cur]);
    if(dp[cur] < dp[root])
        root = cur;
}

void init() {
    for(int i = 1; i <= n; i++) {
        Node[i].clear();
        W[i].clear();
        vis[i] = 0;
    }

    int x, y, z;
    for(int i = 0; i < n - 1; i++) {
        scanf("%d%d%d", &x, &y, &z);
        Node[x].push_back(y);
        Node[y].push_back(x);
        W[x].push_back(z);
        W[y].push_back(z);
    }

    dp[0] = size = n;
    Ans = root = 0;
}

void getdep(int cur, int fa) {
    dep.push_back(d[cur]);
    for(int i = 0; i < Node[cur].size(); i++)
        if(Node[cur][i] != fa && !vis[Node[cur][i]]) {
            d[Node[cur][i]] = d[cur] + W[cur][i];
            getdep(Node[cur][i], cur);
        }
}

int calc(int cur, int Num) {
    dep.clear();
    d[cur] = Num;
    getdep(cur, 0);
    sort(dep.begin(), dep.end());
    int cnt = 0;
    for(int l = 0, r = dep.size() - 1; l < r;)
        if(dep[l] + dep[r] <= k)
            cnt += r - l++;
        else
            r--;
    return cnt;
}
void solve(int cur) {

    Ans += calc(cur, 0);
    vis[cur] = 1;

    for(int i = 0; i < Node[cur].size(); i++) {
        if(!vis[Node[cur][i]]) {
            Ans -= calc(Node[cur][i], W[cur][i]);
            dp[0] = size = Sum[Node[cur][i]];
            dfs(Node[cur][i], root = 0);
            solve(root);
        }
    }
}
int main() {
    while(scanf("%d%d", &n, &k) != EOF && n + k) {
        init();
        dfs(1,0);
        solve(root);
        printf("%d\n", Ans);
    }
    return 0;
}