POJ1741Tree（树的点分治）

题意： 给出一颗树，求树上的两点间最近距离不超过k的点对数量。

思路： 对于一条树路径 只有经过或不经过一个点的情况，对于不经过的情况，把一棵树按这个点拆成好几棵分治就行了，考虑经过这个点的情况，对于这题 可以对这个点延伸出的几棵子树各做一次dfs，记录子树中出现的距离值，对于一棵树的距离值数组，把它排序求一次ans1，再对每棵子树分别求一个自己对自己的ans2，$(ans1-\sum ans2)$即为最后的ans。

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 2e4 + 10;
int n, k, h[N], cnt, rt, sz[N], mx;
int dep[N], bel[N], vis[N], d[N], tot, ans;
struct node {
    int v, w, nt;
} no[N];
void add(int u, int v, int w) {
    no[cnt] = node{v, w, h[u]};
    h[u] = cnt++;
}
void getroot(int u, int fa) {
    sz[u] = 1;
    int ma = 0;
    for(int i = h[u]; ~i; i = no[i].nt) {
        int v = no[i].v;
        if(v != fa && !vis[v]) {
            getroot(v, u);
            sz[u] += sz[v];
            ma = max(ma, sz[v]);
        }
    }
    ma = max(ma, n - sz[u]);
    if(ma < mx) {
        mx = ma;
        rt = u;
    }
}
void getdep(int u, int fa) {
    d[++tot] = dep[u];
    for(int i = h[u]; ~i; i = no[i].nt) {
        int v = no[i].v;
        if(v != fa && !vis[v])
            dep[v] = dep[u] + no[i].w, getdep(v, u);
    }
}
int calc(int u) {
    tot = 0, getdep(u, 0);
    sort(d + 1, d + tot + 1);
    int l = 1, r = tot, sum = 0;
    while(l < r) {
        if(d[l] + d[r] <= k)
            sum += r - l, l++;
        else
            r--;
    }
    return sum;
}
void dfs(int u) {
    dep[u] = 0, vis[u] = 1, ans += calc(u);
    for(int i = h[u]; ~i; i = no[i].nt) {
        int v = no[i].v;
        if(!vis[v])
            mx = 1e9, dep[v] = no[i].w, ans -= calc(v), n = sz[v], rt = 0, getroot(v, 0), dfs(rt);
    }
}
int main() {
    while(~scanf("%d%d", &n, &k) && (n + k)) {
        memset(h, -1, sizeof h);
        cnt = 0, rt = 0, ans = 0;
        memset(vis, 0, sizeof vis);
        for(int u, v, w, i = 1; i < n; i++) {
            scanf("%d%d%d", &u, &v, &w);
            add(u, v, w), add(v, u, w);
        }
        mx = 1e9;
        getroot(1, 0), dfs(rt);
        printf("%d\n", ans);
    }
    return 0;
}