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 = ;

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

int head[N+], 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+], p[N+], siz[N+], g_root, g_min, idx;
bool del[N+];

int getsize(int fa, int u) {
    siz[u] = ;
    for (int i = head[u]; i != -; 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 = ;
    _max = nn - siz[u];
    for (int i = head[u]; i != -; 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 != -; 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 = , j = last - ;
    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 = ;

    g_min = INT_MAX;
    // 得到目前樹的size
    int nn = getsize(-, cur);
    // dp找出root
    // 一旦選了新的root，目前樹的形狀就改變了,siz不再可用
    // 是以每次 dfs 開始前需要重新計算size
    findroot(-, cur, nn);
    root = g_root;
    //删除root，相當于把子樹分割開
    del[root] = ;
    for (int i = head[root]; i != -; i = E[i].nxt) {
        int v = E[i].v;
        if ( !del[v] ) {
            // 分治子樹，累加答案
            ret += dfs(t, v);
            // 每分治完一顆子樹，就将它還原, 相當于重新連回來
            idx = t;
            renew(-, v, E[i].w);
            sort(p + t, p + idx);
            ret -= deal(t, idx);
            t = idx;
        }
    }
    p[t ++] = ;
    sort(p + s, p + t);
    ret += deal(s, t);
    del[root] = ;
    return ret;
}

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