poj1741-Tree（树的分治）

题意：给一棵树，求树上长度小于等于k的链的数量。

题解：http://blog.csdn.net/yang_7_46/article/details/9966455 照着这个博客写的代码。

不到100行，所以不应该算难吧……可是我觉得好难啊……

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;
int n, k;
const int N = ;
struct Edge {
    int to, cost;
};
vector<Edge> g[N];
vector<int> dep;    // 记录一个子树所有节点到根的距离
int sz[N];          // 记录每个子树的大小 最大子树最小的是重心
bool used[N];       // 记录每个节点是否被计算过
int minsz, root;    // 重心的最大子树大小 重心
int size;
int ans;
void up(int &x, int y) { if(y>x) x=y; }
int Scan() {
    int x = ; char C = getchar();
    while (C < '' || C > '') C = getchar();
    while (C >= '' && C <= '') { x = x *  - '' + C, C = getchar(); }
    return x;
}
void getroot(int u, int fa) {               //找到重心
    sz[u] = ;
    int maxn = ;
    for (unsigned i = ; i < g[u].size(); ++i) {
        int v = g[u][i].to;
        if (v == fa || used[v]) continue;   // 因为每次都是找重心 递归求解 求过的不需要再求了
        getroot(v, u);
        sz[u] += sz[v];
        up(maxn, sz[v]);
    }
    up(maxn, size-sz[u]);                   // size不是n 因为每次寻找的树只是一个子树
    if (maxn < minsz) minsz = maxn, root = u;
}
void getdep(int u, int fa, int dis) {       // 寻找子树内每一个结点到根的长度
    dep.push_back(dis);
    sz[u] = ;
    for (unsigned i = ; i < g[u].size(); ++i) {
        int v = g[u][i].to;
        if (v == fa || used[v]) continue;
        getdep(v, u, dis+g[u][i].cost);
        sz[u] += sz[v];
    }
}

int cal(int u, int dis) {                   // dep求得u的所有子树长度，返回的是经过根节点的答案
    dep.clear();
    getdep(u, , dis);
    sort(dep.begin(), dep.end());
    int l = , r = dep.size()-;
    int res = ;
    while (l<r) {
        if (dep[l]+dep[r] <= k) res += r-l++;
        else r--;
    }
    return res;
}

void solve(int u) {                         // 对于每一个结点求解 答案是经过这个结点和不经过这个结点的和
    ans += cal(u, );
    used[u] = true;
    for (unsigned i = ; i < g[u].size(); ++i) {
        int v = g[u][i].to;
        if (used[v]) continue;
        ans -= cal(v, g[u][i].cost);        // 如果两个点位于同一棵子树会重复计算 减去
        minsz = n, root = , size = sz[v];  // size的大小应该是这个子树的大小
        getroot(v, );
        solve(root);
    }
}

int main() {
    while (~scanf("%d%d", &n, &k)) {
        if (n ==  && k == ) break;
        int u, v, c;
        for (int i = ; i <= n; ++i) g[i].clear();
        memset(used, , sizeof used);
        for (int i = ; i < n; ++i) {
            u = Scan(), v = Scan(), c = Scan();
            g[u].push_back(Edge{v, c});
            g[v].push_back(Edge{u, c});
        }
        minsz = n, root = , size = n;
        getroot(, );
        ans = ;
        solve(root);
        printf("%d\n", ans);
    }
    return ;
}