POJ 1741 单次询问树上距离<=K的点对数 点分治

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int MAXN = 1e5 + ;
const int MAXM = 1e5 + ;
int to[MAXM << ], nxt[MAXM << ], Head[MAXN], ed = ;
int cost[MAXM << ];
inline void addedge(int u, int v, int c) {
        to[++ed] = v;
        cost[ed] = c;
        nxt[ed] = Head[u];
        Head[u] = ed;
}
inline void ADD(int u, int v, int c) {
        addedge(u, v, c);
        addedge(v, u, c);
}
int n, k, anser;
int sz[MAXN], f[MAXN], dep[MAXN], sumsz, root;
bool vis[MAXN];
int o[MAXN], cnt;
void getroot(int x, int fa) {
        sz[x] = ;
        f[x] = ;
        for (int i = Head[x]; i; i = nxt[i]) {
                int v = to[i];
                if (v == fa || vis[v]) {
                        continue;
                }
                getroot(v, x);
                sz[x] += sz[v];
                f[x] = max(f[x], sz[v]);
        }
        f[x] = max(f[x], sumsz - sz[x]);
        if (f[x] < f[root]) {
                root = x;
        }
}
void getdeep(int x, int fa) {
        o[++cnt] = dep[x];
        for (int i = Head[x]; i; i = nxt[i]) {
                int v = to[i];
                if (v == fa || vis[v]) {
                        continue;
                }
                dep[v] = dep[x] + cost[i];
                getdeep(v, x);
        }
}
int calc(int x, int d) {
        cnt = ;
        dep[x] = d;
        getdeep(x, );
        sort(o + , o + cnt + );
        int l = , r = cnt, ansnow = ;
        while (l < r) {
                if (o[l] + o[r] <= k) {
                        ansnow += r - l, l++;
                } else {
                        r--;
                }
        }
        return ansnow;
}
void solve(int x) {
        anser += calc(x, ); //满足Depth(i)+Depth(j)<=k的对数
        vis[x] = ;
        int totsz = sumsz;
        for (int i = Head[x]; i; i = nxt[i]) {
                int v = to[i];
                if (vis[v]) {
                        continue;
                }
                anser -= calc(v, cost[i]);//满足Depth(i)+Depth(j)<=k且Belong(i)=Belong(j)的对数
                root = ;
                sumsz = sz[v] > sz[x] ? totsz - sz[x] : sz[v];
                getroot(v, );
                solve(root);
        }
}
int main() {
        while (scanf("%d %d", &n, &k) && (n + k)) {
                cnt = anser = ;
                memset(Head, , sizeof(Head));
                memset(vis, , sizeof(vis));
                ed = ;
                int u, v, c;
                for (int i = ; i < n; i++) {
                        scanf("%d %d %d", &u, &v, &c);
                        ADD(u, v, c);
                }
                root = , sumsz = f[] = n;
                getroot(, );
                solve(root);
                printf("%d\n", anser);
        }
        return ;
}