POJ 1741 Tree （树分治入门）

参考:
分治算法在树的路径问题中的应用 - 漆子超
http://www.cnblogs.com/kuangbin/p/3454883.html
http://www.cnblogs.com/staginner/archive/2012/07/04/2576720.html

我写的是点分治，每次 solve 过程由调用者负责分解并还原子树。
dp 求 root 的时候，一定要用之前一步 dfs 得到的 tree size。
最开始没理解透，直接用初始的n，debug很久。。

const int N = 10000;

struct Edge {
    int v, w, nxt;
} E[N * 2 + 5];

int head[N+5], enm;

void add_edge(int u, int v, int w) {
    E[enm] = (Edge) {v, w, head[u]};
    head[u] = enm ++;
}

int n, k, fa[N+5], p[N+5], siz[N+5], g_root, g_min, idx;
bool del[N+5];

int getsize(int fa, int u) {
    siz[u] = 1;
    for (int i = head[u]; i != -1; i = E[i].nxt) {
        int v = E[i].v;
        if ( !del[v] && v != fa ) {
            siz[u] += getsize(u, v);
        }
    }
    return siz[u];
}

void findroot(int fa, int u, int nn) {
    int _max = 0;
    _max = nn - siz[u];
    for (int i = head[u]; i != -1; i = E[i].nxt) {
        int v = E[i].v;
        if ( !del[v] && v != fa ) {
            findroot(u, v, nn);
            if ( siz[v] > _max ) {
                _max = siz[v];
            }
        }
    }
    if ( _max < g_min ) {
        g_min = _max;
        g_root = u;
    }
}

void renew(int fa, int u, int d) {
    p[idx ++] = d;
    for (int i = head[u]; i != -1; i = E[i].nxt) {
        int v = E[i].v;
        if ( !del[v] && v != fa ) renew(u, v, d + E[i].w);
    }
}

// calculate how many p[i] + p[j] <= k in p[first, last), i <= j
int deal(int first, int last) {
    int ret = 0, j = last - 1;
    for (int i = first; i < last; ++ i) {
        if ( p[i] > k ) break;
        while ( i < j && p[i] + p[j] > k ) -- j;
        if ( i < j )
            ret += j - i;
        else
            break;
    }
    return ret;
}

int dfs(int s, int cur) {
    int root, tmp, t = s, ret = 0;

    g_min = INT_MAX;
    // 得到当前树的size
    int nn = getsize(-1, cur);
    // dp找出root
    // 一旦选了新的root，当前树的形状就改变了,siz不再可用
    // 所以每次 dfs 开始前需要重新计算size
    findroot(-1, cur, nn);
    root = g_root;
    //删除root，相当于把子树分割开
    del[root] = 1;
    for (int i = head[root]; i != -1; i = E[i].nxt) {
        int v = E[i].v;
        if ( !del[v] ) {
            // 分治子树，累加答案
            ret += dfs(t, v);
            // 每分治完一颗子树，就将它还原, 相当于重新连回来
            idx = t;
            renew(-1, v, E[i].w);
            sort(p + t, p + idx);
            ret -= deal(t, idx);
            t = idx;
        }
    }
    p[t ++] = 0;
    sort(p + s, p + t);
    ret += deal(s, t);
    del[root] = 0;
    return ret;
}

int main() {
    while ( scanf("%d%d", &n, &k) != EOF && n + k ) {
        memset(head, -1, sizeof(head));
        enm = 0;
        rep(i, 1, n-1) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            add_edge(u, v, w);
            add_edge(v, u, w);
        }
        int ans = dfs(0, 1);
        printf("%d\n", ans);
    }
    return 0;
}