POJ1741 tree

#include <cstdio>
#include <cstring>

#include <algorithm>

#define LOG(FMT...) // fprintf(stderr, "[LOG]: "FMT)

using namespace std;

typedef long long ll;

struct edge {
    int v, w;
    edge* next;
};

const int N = 10010;

int n, k, ocnt, ql, qr;
ll ans;
edge* ptop;

bool vis[N], cut[N];
int ord[N], q[N], sum[N], mex[N], dis[N], d[N];
edge pool[N * 2];
edge* g[N];

void add_edge(int u, int v, int w);
void calc_order(int u);
ll calc(int u);

int main() {
    mex[0] = 1 << 20;
    while (scanf("%d%d", &n, &k), n != 0) {
        ans = 0;
        ptop = pool;
        ocnt = 0;
        memset(g, 0, sizeof(g));
        for (int rep = 1; rep < n; ++rep) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            add_edge(u, v, w);
            add_edge(v, u, w);
        }
        memset(cut, 0, sizeof(cut));
        calc_order(1);
        LOG("FINE ORDER\n");
        memset(cut, 0, sizeof(cut));
        for (int i = 0; i < n; ++i) {
            int u = ord[i];
            LOG("ORDER%d: %d\n", i, u);
            ql = 0;
            qr = 1;
            q[0] = u;
            dis[u] = 0;
            while (ql < qr) {
                int v = q[ql++];
                vis[v] = true;
                for (edge* p = g[v]; p; p = p->next)
                    if (!vis[p->v] && !cut[p->v]) {
                        dis[p->v] = dis[v] + p->w;
                        q[qr++] = p->v;
                    }
            }
            for (int i = 0; i < qr; ++i)
                vis[q[i]] = false;
            ans += calc(u);
            LOG("ADDED\n");
            cut[u] = true;
            for (edge* p = g[u]; p; p = p->next)
                if (!cut[p->v])
                    ans -= calc(p->v);
            LOG("%lld\n", ans);
        }
        printf("%lld\n", (ans - n) / 2);
    }
    return 0;
}

ll calc(int u) {
    ll ret = 0;
    ql = 0;
    qr = 1;
    q[0] = u;
    while (ql < qr) {
        u = q[ql];
        d[ql] = dis[u];
        ++ql;
        vis[u] = true;
        for (edge* p = g[u]; p; p = p->next)
            if (!vis[p->v] && !cut[p->v])
                q[qr++] = p->v;
    }
    for (int i = 0; i < qr; ++i)
        vis[q[i]] = false;
    sort(d, d + qr);
    int lim = qr;
    LOG("[LIM]%d\n", lim);
    for (ql = 0, --qr; qr >= 0; --qr) {
        while (d[ql] + d[qr] <= k && ql < lim)
            ++ql;
        LOG("%d %d\n", d[ql - 1], d[qr]);
        ret += ql;
    }
    return ret;
}

void add_edge(int u, int v, int w) {
    edge* p = ptop;
    ++ptop;
    p->v = v;
    p->w = w;
    p->next = g[u];
    g[u] = p;
}

void calc_order(int u) {
    ql = 0;
    qr = 1;
    q[0] = u;
    while (ql < qr) {
        u = q[ql++];
        vis[u] = true;
        for (edge* p = g[u]; p; p = p->next)
            if (!vis[p->v] && !cut[p->v])
                q[qr++] = p->v;
    }
    u = 0;
    for (int i = qr - 1; i >= 0; --i) {
        int v = q[i];
        vis[v] = false;
        sum[v] = 1;
        mex[v] = 0;
        for (edge* p = g[v]; p; p = p->next)
            if (!vis[p->v] && !cut[p->v]) {
                sum[v] += sum[p->v];
                mex[v] = max(mex[v], sum[p->v]);
            }
        mex[v] = max(mex[v], qr - sum[v]);
        if (mex[v] < mex[u])
            u = v;
    }
    ord[ocnt++] = u;
    cut[u] = true;
    for (edge* p = g[u]; p; p = p->next)
        if (!cut[p->v])
            calc_order(p->v);
}