*POJ 1741 - Tree(树形DP‘树的分治)

题目：

http://poj.org/problem?id=1741

题意：

求树中两点之间的距离小于等于k的对数。

思路：

树的点分治。看论文+模板。

AC.

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;
const int maxn = 10005;
int n, k;
int vis[maxn], ans, root, num;

int tot, head[maxn];
struct Edge {
    int to, w, next;
}edge[maxn*5];

void addedge(int u, int v, int w)
{
    edge[tot].to = v;
    edge[tot].w = w;
    edge[tot].next = head[u];
    head[u] = tot++;
}
int mx[maxn], size[maxn], mi, dis[maxn];
void dfssize(int u, int fa) //处理子树的大小
{
    size[u] = 1;
    mx[u] = 0;
    for(int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if(v != fa && !vis[v]) {
            dfssize(v, u);
            size[u] += size[v];
            if(size[v] > mx[u]) mx[u] = size[v];

        }
    }
}
void dfsroot(int r, int u, int fa) //求重心
{
    if(size[r] - size[u] > mx[u]) mx[u] = size[r] - size[u];
    if(mx[u] < mi) mi = mx[u], root = u;
    for(int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if(v != fa && !vis[v]) dfsroot(r, v, u);
    }
}
void dfsdis(int u, int d, int fa) //求距离
{
    dis[num++] = d;
    for(int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if(v != fa && !vis[v]) dfsdis(v, d+edge[i].w, u);
    }
}
int cal(int u, int d)
{
    int ret = 0;
    num = 0;
    dfsdis(u, d, 0);
    sort(dis, dis+num);
    int i = 0, j = num-1;
    while(i < j) {
        while(dis[i] + dis[j] > k && i < j) j--;
        ret += j-i;
        i++;
    }
    return ret;
}
void dfs(int u)
{
    mi = n;
    dfssize(u, 0);
    dfsroot(u, u, 0);
    ans += cal(root, 0);
    vis[root] = 1;
    for(int i = head[root]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if(!vis[v]) {
            ans -= cal(v, edge[i].w);
            dfs(v);
        }
    }
}
void init()
{
    ans = tot = 0;
    memset(head, -1, sizeof(head));
    memset(vis, 0, sizeof(vis));
}
int main()
{
    //freopen("in", "r", stdin);
    while(~scanf("%d %d", &n, &k)) {
        if(n == 0 && k == 0) break;

        init();
        int u, v, w;
        for(int i = 1; i < n; ++i) {
            scanf("%d %d %d", &u, &v, &w);
            addedge(u, v, w);
            addedge(v, u, w);
        }

        dfs(1);
        printf("%d\n", ans);
    }
    return 0;
}