Tree

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001). 
Define dist(u,v)=The min distance between node u and v. 
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k. 
Write a program that will count how many pairs which are valid for a given tree. 

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l. 
The last test case is followed by two zeros. 

Output

For each test case output the answer on a single line.

Sample Input

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

Sample Output

8

Source

[email protected]

題目大意：n個點的樹，每條邊有權重w，求有多少個點對(x, y)滿足x到y的距離小於等於k

最基礎的想法，對每個點都遍歷一遍整個樹，得到合法點對的個數，複雜度O(n2)無法滿足需求。
注意到，以root為根的樹上的路徑，要麼經過root，要麼在root的子樹內部，對於子樹內部的路徑可以通過遞迴得到解決；
對於經過root的路徑，比如root有兩個子樹x和y，用depth[i]表示root到i節點的距離，那麼合法的路徑即為滿足depth[i] + depth[j] <= k
(i∈x，j∈y)的個數，進一步的可以轉換為 i∈(x, y)，j∈(x, y)的合法點對 - i∈x，j∈x的合法點對 - i∈y，j∈y的合法點對
對於求depth[i] + depth[j] <= k 可以通過排序後雙指標線性的得到答案，因此對於深度為D的樹，每一層的複雜度為O(nlogn)，所以總的複雜度為O(D*nlogn)
為了防止D退化為n，每次找root應為樹的重心，求樹的重心可以線性得到，因此總的複雜度為O(nlognlogn)

#include <cstdio>
#include <vector>
#include <cstring>
#include <utility>
#include <iostream>
#include <algorithm>
using namespace std;

const int maxn = 1e4 + 10;
vector<pair<int, int> >G[maxn];
vector<int> vec;
int vis[maxn], sz[maxn], maxson[maxn], depth[maxn];
int n, k;
int ans, root, allnode;

void find_root(int u, int fa) {
    sz[u] = 1;
    maxson[u] = 0; //maxson[u]表示去除節點u之後節點數最多的樹的節點數
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i].first;
        int w = G[u][i].second;
        if (v == fa || vis[v]) continue;
        find_root(v, u);
        maxson[u] = max(sz[v], maxson[u]);  //用子樹v的節點數更新maxson[u]
        sz[u] += sz[v];
    }
    maxson[u] = max(maxson[u], allnode - sz[u]); 
    if (maxson[u] < maxson[root]) root = u;
}

//計運算元樹u下所有節點的depth
void dfs(int u, int fa) { 
    vec.push_back(depth[u]);
    sz[u] = 1;
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i].first;
        int w = G[u][i].second;
        if (v == fa || vis[v]) continue;
        depth[v] = depth[u] + w;
        dfs(v, u);
        sz[u] += sz[v];
    }
}

int calc(int u, int fa, int init) {
    vec.clear(); //節點u下的所有節點的depth會存在vec中
    depth[u] = init;
    dfs(u, fa);
    sort(vec.begin(), vec.end()); //對vec排序後用雙指標掃一遍得到合法點對數
    int cnt = 0, l = 0, r = vec.size() - 1;
    for (; l < r;) {
        if (vec[l] + vec[r] <= k) {
            cnt += (r - l);
            l++;
        }
        else r--;
    }
    return cnt;
}

void solve(int u, int fa) {
    vis[u] = 1;
    ans += calc(u, fa, 0); //計算當前根下所有合法點對的個數
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i].first;
        int w = G[u][i].second;
        if (v == fa || vis[v]) continue;
        ans -= calc(v, u, w); //去除同一子樹內部的合法點對
        root = 0;
        allnode = sz[v];
        find_root(v, u);
        solve(root, u);
    }
}

int main() {
    while (scanf("%d%d", &n, &k) != EOF) {
        if (n == 0 && k == 0) break;
        for (int i = 1; i <= n; i++) G[i].clear();
        ans = 0;
        maxson[0] = n + 1;
        memset(vis, 0, sizeof(vis));
        for (int i = 1; i < n; i++) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            G[u].push_back({v, w});
            G[v].push_back({u, w});
        }
        root = 0;
        allnode = n;
        find_root(1, -1);
        solve(root, -1);
        printf("%d\n", ans);
    }
    return 0;
}