POJ 1741 Tree 树分治

题意：

给出一颗有\(n (n \leq 10^4)\)个节点的树，和一个\(k\)。统计有多少个点对\(u, \, v(u \neq v)\)满足\(u\)到\(v\)的最短距离不超过\(k\)。

分析：

树分治的入门题，可以参考论文《分治算法在树的路径问题中的应用》。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <map>
using namespace std;

typedef pair<int, int> PII;

const int maxn = 10000 + 10;
const int INF = 0x3f3f3f3f;

struct Edge
{
    int v, w, nxt;
    Edge() {}
    Edge(int v, int w, int nxt): v(v), w(w), nxt(nxt) {}
};

int n, k;

int ecnt, head[maxn];
Edge edges[maxn * 2];

void AddEdge(int u, int v, int w) {
    edges[ecnt] = Edge(v, w, head[u]);
    head[u] = ecnt++;
}

int sz[maxn], fa[maxn];  //sz[u]是以u为根子树的大小，fa[u]是u的父亲
bool del[maxn];  //del[u]表示u作为某颗子树的重心删除的标记
int ans, mins, centroid;  //最终答案，最小的最大子树以及重心
vector<int> d, d2;

//计算sz和fa
void dfs(int u) {
    sz[u] = 1;
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(v == fa[u] || del[v]) continue;
        fa[v] = u;
        dfs(v);
        sz[u] += sz[v];
    }
}

//计算重心
void dfs2(int u, int t) {
    int m = 0;
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(v == fa[u] || del[v]) continue;
        m = max(m, sz[v]);
        dfs2(v, t);
    }
    m = max(m, t - sz[u]);
    if(m < mins) { mins = m; centroid = u; }
}

//统计所有点到根节点的距离
void getdist(int u, int p, int dist, vector<int>& d) {
    d.push_back(dist);
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v, w = edges[i].w;
        if(v == p || del[v]) continue;
        getdist(v, u, dist + w, d);
    }
}

//统计符合要求点对个数
int cntpair(vector<int>& d) {
    int ans = 0;
    sort(d.begin(), d.end());
    int j = d.size();
    for(int i = 0; i < d.size(); i++) {
        while(j > 0 && d[i] + d[j-1] > k) j--;
        ans += j - (j > i ? 1 : 0);  //去掉和自己成为一对的情况
    }
    return ans / 2;
}

//分治过程
void solve(int u) {
    fa[u] = 0;
    dfs(u);
    mins = INF;
    dfs2(u, sz[u]);
    int s = centroid;
    del[s] = true;

    for(int i = head[s]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(del[v]) continue;
        solve(v);
    }

    d.clear();
    d.push_back(0);
    for(int i = head[s]; ~i; i = edges[i].nxt) {
        int v = edges[i].v, w = edges[i].w;
        if(del[v]) continue;
        d2.clear();
        getdist(v, s, w, d2);
        ans -= cntpair(d2);  //去掉计重部分
        d.insert(d.end(), d2.begin(), d2.end());
    }
    ans += cntpair(d);
    del[s] = false;
}

int main()
{
    while(scanf("%d%d", &n, &k) == 2) {
        if(!n && !k) break;
        ecnt = 0;
        memset(head, -1, sizeof(head));
        for(int i = 1; i < n; i++) {
            int u, v, w; scanf("%d%d%d", &u, &v, &w);
            AddEdge(u, v, w);
            AddEdge(v, u, w);
        }
        
        ans = 0;
        solve(1);
        printf("%d\n", ans);
    }

    return 0;
}