POJ：1741-Tree(树的点分治模板)

Tree

Time Limit: 1000MS Memory Limit: 30000K
Total Submissions: 31615 Accepted: 10567

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

Source
LouTiancheng@POJ

这题就是点分治的模板题，给你一个树，树的边有权值，问任意路径权值和不大于ｋ的有多少条。其实思路还是很简单的，归并排序的思路，不过把线性的数组放到树上面了，这就需要找重心等一系列操作。《挑战程序设计》上面讲得很清楚了。

#include <stdio.h>
#include <cstring>
#include <algorithm>
#include <vector>
#include <climits>
using namespace std;
const int maxn = 1e4+100;

int n, k, subtree_size[maxn], ans;
bool centroid[maxn];

struct edge {
    int to, len;

    edge(int To, int Len):
        to(To), len(Len){}
};
vector <edge> ve[maxn];

void init() {
    for(int i=0;i<=n;i++) {
        ve[i].clear();
    }
    memset(centroid, 0, sizeof (centroid));
    for(int i=1;i<n;i++) {
        int a, b, len;
        scanf("%d%d%d", &a, &b, &len);
        a--, b--;
        ve[a].push_back(edge(b, len));
        ve[b].push_back(edge(a, len));
    }
}

//计算以ｖ为根节点的树的节点数目
int compute_subtree_size(int v, int p) {
    int c = 1;
    for(int i=0;i<ve[v].size();i++) {
        int to = ve[v][i].to;
        if(to == p || centroid[to]) continue;
        c += compute_subtree_size(to, v);
    }
    subtree_size[v] = c;
    return c;
}

//找到以ｖ为根节点的pair<树的重心的最大子树尺寸,重心所在的节点>
pair <int, int> search_centroid(int v, int p, int t) {
    pair <int, int> res = make_pair(INT_MAX, -1);
    int s = 1, m = 0;

    for(int i=0;i<ve[v].size();i++) {
        int w = ve[v][i].to;

        if(w == p || centroid[w]) continue;

        res = min(res, search_centroid(w, v, t));
        m = max(m, subtree_size[w]);
        s += subtree_size[w];
    }

    m = max(m, t-s);
    res = min(res, make_pair(m, v));

    return res;
}

//找到ｖ为根节点的树的所有路径长度
void enumarate_paths(int v, int p, int d, vector<int> &ds) {
    ds.push_back(d);
    for(int i=0;i<ve[v].size();i++) {
        int w = ve[v][i].to;

        if(w == p || centroid[w]) continue;

        enumarate_paths(w, v, d+ve[v][i].len, ds);
    }
}

int count_pairs(vector <int> &ds) {
    int ans = 0;
    sort(ds.begin(), ds.end());

    int j = ds.size();
    for(int i=0;i<ds.size();i++) {
        while(j > 0 && ds[j-1] + ds[i] > k) j--;
        ans += j - (j > i ? 1 : 0);
    }

    return ans/2;
}

//划分子树并且找到相关的值
void solve_subproblem(int v) {
    compute_subtree_size(v, -1);
    int s = search_centroid(v, -1, subtree_size[v]).second;
    centroid[s] = true;

    for(int i=0;i<ve[s].size();i++) {
        if(centroid[ve[s][i].to]) continue;
        solve_subproblem(ve[s][i].to);
    }

    vector <int> ds;
    ds.clear();
    ds.push_back(0);
    for(int i=0;i<ve[s].size();i++) {
        int w = ve[s][i].to;
        if(centroid[w]) continue;

        vector <int> tds;
        tds.clear();
        enumarate_paths(w, s, ve[s][i].len, tds);

        ans -= count_pairs(tds);

        ds.insert(ds.end(), tds.begin(), tds.end());
    }

    ans += count_pairs(ds);
    centroid[s] = false;
}

void solve() {
    ans = 0;
    solve_subproblem(0);
    printf("%d\n", ans);
}

int main() {
    //freopen("1.in", "r", stdin);
    while(scanf("%d%d", &n, &k) && n+k) {
        init();
        solve();
    }
    return 0;
}