POJ1741Tree(树形dp/树分治)

题目大意：给一棵边带权树，问两点之间的距离小于等于K的点对有多少个。

具体题解看论文。

#include<iostream>
#include<algorithm>
#include<cstring>
#include<cstdio>
using namespace std;
const int maxn = 21000;
struct Node
{
    int to, w;
    int sum, maxn;
    Node *nxt;
}tree[maxn], *head[maxn], tp[maxn];
bool vis[maxn];
int dist[maxn], Size[maxn], sign[maxn];
int ans, ptr, root, n, k, tot;
//ans为所求答案,ptr边的计数
void AddEdge(int u, int v, int w)
{
    tree[ptr].to = v;
    tree[ptr].w = w;
    tree[ptr].nxt = head[u];
    head[u] = &tree[ptr++];
}
void Dfs(int x, int past)//求树的重心
{
    tp[x].sum = tp[x].maxn = 0;
    Node *p = head[x];
    while(p != NULL)
    {
        if(p->to != past && vis[p->to] == false)
        {
            Dfs(p->to, x);
            tp[x].sum += tp[p->to].sum;//以该节点为根节点的树的节点个数
            tp[x].maxn = max(tp[x].maxn, tp[p->to].sum);//最大子树
        }
        p = p->nxt;
    }
    tp[x].sum++;//加上根节点
    sign[tot] = x;//记录下这个节点
    Size[tot++] = tp[x].maxn;//该节点最大子树的节点个数
}
int GetRoot(int x)
{
    tot = 0;
    Dfs(x, 0);
    int maxn = 0x7f7f7f7f, maxsign = 0, cnt = tp[x].sum;
    for(int i = 0; i < tot; i++)
    {
        Size[i] = max(Size[i], cnt - Size[i]-1);//获得当前节点的最大子树
        if(Size[i] < maxn)
        {
            maxn = Size[i];
            maxsign = sign[i];
        }
    }
    return maxsign;
}
void CalcDist(int s, int past, int dis)//计算各个点到当前根节点的距离
{
    Node *p = head[s];
    dist[tot++] = dis;
    while(p != NULL)
    {
        if(p->to != past && vis[p->to] == false && dis + p->w <= k)
            CalcDist(p->to, s, dis+p->w);
        p = p->nxt;
    }
}
void CountSum(int x)
{
    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 CountDel(int x)
{
    vis[x] = true;//已经处理完以当前节点为根节点的数,标上记号
    Node *p = head[x];
    while(p != NULL)
    {
        if(vis[p->to] == false)
        {
            tot = 0;
            CalcDist(p->to, x, p->w);
            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->nxt;
    }
}
void Deal(int s, int past)//s当前点,past父节点
{
    root = GetRoot(s);//获得当前这棵数的根节点
    tot = 0;//记录当前树中节点个数
    CalcDist(root, 0, 0);//计算各个点到根节点的距离
    CountSum(root);//计算depth(i)+depth(j) <= k的对数
    CountDel(root);//计算同一个子树中depth(i)+depth(j) <= k的对数，即belong(i)==belong(j)
    Node *p = head[root];
    while(p != NULL)//分治同样的方法处理子树
    {
        if(p->to != past && vis[p->to] == false)
            Deal(p->to, root);
        p = p->nxt;
    }
}
int main()
{
    while(scanf("%d%d", &n, &k) != EOF && n + k != 0)
    {
        ans = ptr = 0;
        memset(vis, false, sizeof(vis));
        memset(head, 0, sizeof(head));
//        for(int i = 0; i < maxn; i++)
//            vis[i] = false, head[i] = NULL;
        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);
        }
        Deal(1, 0);
        printf("%d\n", ans);
    }
    return 0;
}