poj 1741 Tree(树形DP+分治)难

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

2、题目大意：

给一棵树，dist(u,v)是指的uv两个顶点之间的最短距离，现在给出一个k值，

求在这棵树上有多少对(u,v)满足dist(u,v)<=k

如果uv在同一个分支上，那么dist(u,v)就死uv间的距离，如果uv不在同一个分支上，那么他们之间的最短距离就等于他们到最近的公共父节点的距离之和。

这道题目如果用dfs搜索每一个根节点得用n*n的时间，会超时，看http://blog.csdn.net/woshi250hua/article/details/7723400介绍的需要用到分治，

3、题目

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 8667		Accepted: 2591

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.
The last test case is followed by two zeros.

Output

For each test case output the answer on a single line.

Sample Input

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

Sample Output

8

AC代码参考http://blog.csdn.net/woshi250hua/article/details/7723400

#include <stdio.h>
#include <string.h>
#include <algorithm>
using namespace std;
#define MAX 21000
#define INF 2147483647
#define max(a,b) (a)>(b)?(a):(b)
#define min(a,b) (a)<(b)?(a):(b)

struct node {
    int v, len; //v表示邻接点，len表示边权
    int sum, bal; //sum表示子孙节点个数，bal表示拆当前点后子树的最大结点数
    node *next;
} tree[MAX], *head[MAX], tp[MAX];
int n, ptr, ans, root;
int tot, k, dist[MAX], vis[MAX];
int size[MAX], sign[MAX]; //size表示最大分支的结点数,sign是一个hash数组

void Initial() {

    ans = ptr = 0;
    for (int i = 0; i < MAX; ++i)
        vis[i] = 0,head[i] = NULL;
}

void AddEdge(int a, int b, int c) {

    tree[ptr].v = b, tree[ptr].len = c;
    tree[ptr].next = head[a], head[a] = &tree[ptr++];
}

void Dfs(int s, int pa) {

    tp[s].sum = tp[s].bal = 0;
    node *p = head[s];
    while (p != NULL) {

        if (p->v != pa && vis[p->v] == 0) {//在哪赋值成1？？？

            Dfs(p->v, s);
            tp[s].sum += tp[p->v].sum;                //累计子节点个数
            tp[s].bal = max(tp[s].bal, tp[p->v].sum); //找最大分支
        }
        p = p->next;
    }
    tp[s].sum++;             //自己
    sign[tot] = s;           //hash
    size[tot++] = tp[s].bal; //记录每个最大分支的结点数
}

int GetRoot(int s) {
    tot = 0, Dfs(s, 0);
    int maxx = INF, maxi, cnt = tp[s].sum;
    for (int i = 0; i < tot; ++i) {
        size[i] = max(size[i], cnt - size[i]);
        if (size[i] < maxx) {
            maxx = size[i];
            maxi = sign[i];
        }
    }
    return maxi;
}

void GetDist(int s, int pa, int dis) {
    //保存每个结点到根节点的距离
    node *p = head[s];
    dist[tot++] = dis;


    while (p != NULL) {

        if (p->v != pa && vis[p->v] == 0 && dis + p->len <= k)
            GetDist(p->v, s, dis + p->len);
        p = p->next;
    }
}

void Count1(int s) {

    sort(dist, dist + tot);
    int left = 0, right = tot - 1;


    while (left < right) {

        if (dist[left] + dist[right] <= k)
            ans += right - left, left++;
        else right--;
    }
}
void Count2(int s) {

    vis[s] = 1;
    node *p = head[s];


    while (p != NULL) {

        if (vis[p->v] == 0) {

            tot = 0, GetDist(p->v, s, p->len);
            sort(dist, dist + tot);
            int left = 0, right = tot - 1;
            while (left < right) {

                if (dist[left] + dist[right] <= k)
                    ans -= right - left, left++;
                else right--;
            }
        }
        p = p->next;
    }
}

void Solve(int s, int pa) {

    root = GetRoot(s);
    tot = 0,GetDist(root, 0, 0);
    Count1(root);
    Count2(root);  // ans += count1 - coutn2;


    node *p = head[root];
    while (p != NULL) {

        if (p->v != pa && vis[p->v] == 0)
            Solve(p->v, root);
        p = p->next;
    }
}


int main()
{
    int i, j, a, b, c;


    while (scanf("%d%d", &n, &k),n + k) {

        Initial();
        for (i = 1; i < n; ++i) {

            scanf("%d%d%d", &a, &b, &c);
            AddEdge(a, b, c);
            AddEdge(b, a, c);
        }


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

}