poj1741 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\leq 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

Solution

点分治裸题。简介见点分治简介。注意要去重（见 $calcPair$ ）。

#include<bits/stdc++.h>
using namespace std;

#define N 100001
#define rep(i, a, b) for (int i = a; i <= b; i++)
#define drp(i, a, b) for (int i = a; i >= b; i--)
#define fech(i, x) for (int i = 0; i < x.size(); i++)
#define pii pair<int, int>
#define INF 0x7fffffff

inline int read() {
    int x = 0, flag = 1; char ch = getchar(); while (!isdigit(ch)) { if (!(ch ^ '-')) flag = -1; ch = getchar(); }
    while (isdigit(ch)) x = (x << 1) + (x << 3) + ch - '0', ch = getchar(); return x * flag;
}

inline void write(int x) {
    if (!x) { putchar('0'); return; } if (x < 0) putchar('-'), x = -x;
    char buf[20] = ""; int top = 0; while (x) buf[++top] = x % 10 + '0', x /= 10; while (top) putchar(buf[top--]);
}

int n, K;
struct edgeType { int u, v, w; }eg[N]; int tot;
vector<int> g[N];
bool centr[N];
int Size[N];
int ans;

pii calcCentr(int u, int f, int sz) {
    int mx = 0, sum = 1; pii res = pii{ INF, -1 };
    fech(i, g[u]) {
        edgeType e = eg[g[u][i]]; if (!(e.v ^ f) || centr[e.v]) continue;
        res = min(res, calcCentr(e.v, u, sz)); mx = max(mx, Size[e.v]); sum += Size[e.v];
    }
    mx = max(mx, sz - sum); res = min(res, pii{ mx, u }); return res;
}

int calcSize(int u, int f) {
    Size[u] = 1;
    fech(i, g[u]) {
        edgeType e = eg[g[u][i]]; if (!(e.v ^ f) || centr[e.v]) continue;
        calcSize(e.v, u); Size[u] += Size[e.v];
    }
    return Size[u];
}

void calcDis(int u, int f, int d, vector<int>& ds) {
    ds.push_back(d);
    fech(i, g[u]) {
        edgeType e = eg[g[u][i]]; if (!(e.v ^ f) || centr[e.v]) continue;
        calcDis(e.v, u, e.w + d, ds);
    }
}

inline int calcPair(vector<int>& ds) {
    int res = 0;
    sort(ds.begin(), ds.end());
    int j = ds.size();
    fech(i, ds) {
        while (j > 0 && ds[i] + ds[j - 1] > K) j--;
        res += j - (j > i ? 1 : 0);
    }
    return res / 2;
}

void solve(int u) {
    calcSize(u, 0);
    int s = calcCentr(u, 0, Size[u]).second;
    centr[s] = 1;
    fech(i, g[u]) if (!centr[eg[g[u][i]].v]) solve(eg[g[u][i]].v);

    vector<int> ds; ds.push_back(0);
    fech(i, g[u]) {
        edgeType e = eg[g[u][i]]; if (centr[e.v]) continue;
        vector<int> tds; calcDis(e.v, u, e.w, tds);
        ans -= calcPair(tds);
        ds.insert(ds.begin(), tds.begin(), tds.end());
    }
    ans += calcPair(ds);
    centr[s] = 0;
}

int main() {
    while (scanf("%d%d", &n, &K) == 2 && n && K) {
        tot = 0; rep(i, 1, n) g[i].clear();
        rep(i, 2, n) {
            int u = read(), v = read(), w = read();
            eg[++tot] = edgeType{ u, v, w }; g[u].push_back(tot);
            eg[++tot] = edgeType{ v, u, w }; g[v].push_back(tot);
        }
        ans = 0; solve(1); write(ans); puts("");
    }
    return 0;
}