POJ 2114 Boatherds 树的分治

题目：http://poj.org/problem?id=2114

题意：给一个有n个点的树，树有边权。问有没有一个点对，使两点间路径边权和为一个特定值，有的话输出AYE，没有的话输出NYE。输入格式如下：首先输入一个n，接下来有n行，第j行的数据中，形式如ai bi，代表j点和ai点之间有边，权值为bi。每行的输入直到遇到ai为0，则这一行结束。然后每次输入一个数，询问树中有没有两点间路径边权和等于这个数

思路：树的分治，有点类似于poj 1741，可以更改一下统计函数。注意判断时不能直接找到一个值等于输入的询问就认为有这样的区间，可能是同一个子树里面的两个节点

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int N = 10010;
struct edge
{
    int to, val, next;
}g[N*2];
int head[N], siz[N], sonsiz[N], dis[N];
int n, cnt, num, tmp, root, res, k;
bool vis[N];
void add_edge(int v, int u, int val)
{
    g[cnt].to = u;
    g[cnt].val = val;
    g[cnt].next = head[v];
    head[v] = cnt++;
}
void dfs_siz(int v, int fa)
{
    siz[v] = 1, sonsiz[v] = 0;
    for(int i = head[v]; i != -1; i = g[i].next)
    {
        int u = g[i].to;
        if(u != fa && !vis[u])
        {
            dfs_siz(u, v);
            siz[v] += siz[u];
            sonsiz[v] = max(sonsiz[v], siz[u]);
        }
    }
}
void dfs_root(int r, int v, int fa)
{
    sonsiz[v] = max(sonsiz[v], siz[r] - siz[v]);
    if(sonsiz[v] < tmp) tmp = sonsiz[v], root = v;
    for(int i = head[v]; i != -1; i = g[i].next)
    {
        int u = g[i].to;
        if(u != fa && !vis[u]) dfs_root(r, u, v);
    }
}
void dfs_dis(int v, int d, int fa)
{
    dis[num++] = d;
    for(int i = head[v]; i != -1; i = g[i].next)
    {
        int u = g[i].to;
        if(u != fa && !vis[u]) dfs_dis(u, d + g[i].val, v);
    }
}
int calc(int v, int d)
{
    num = 0;
    int ans = 0;
    dfs_dis(v, d, 0);
    sort(dis, dis + num);
    int i = 0, j = num - 1;
    while(i < j)
    {
        if(dis[i] + dis[j] < k) i++;
        else if(dis[i] + dis[j] > k) j--;
        else
        {
            if(dis[i] == dis[j])
            {
                ans += (j - i + 1) * (j - i) / 2; break;
            }
            else
            {
                int l = i, r = j;
                while(dis[l] == dis[i]) l++;
                while(dis[r] == dis[j]) r--;
                ans += (l - i) * (j - r);
                i = l, j = r;
            }
        }
    }
    return ans;
}
void dfs(int v)
{
    tmp = n;
    dfs_siz(v, 0);
    dfs_root(v, v, 0);
    res += calc(root, 0);
    vis[root] = 1;
    for(int i = head[root]; i != -1; i = g[i].next)
    {
        int u = g[i].to;
        if(!vis[u])
        {
            res -= calc(u, g[i].val);
            dfs(u);
        }
    }
}
int main()
{
    int a, b;
    while(scanf("%d", &n), n)
    {
        cnt = 0;
        memset(head, -1, sizeof head);
        for(int i = 1; i <= n; i++)
        {
            while(true)
            {
                scanf("%d", &a);
                if(! a) break;
                scanf("%d", &b);
                add_edge(i, a, b);
                add_edge(a, i, b);
            }
        }
        while(scanf("%d", &k), k)
        {
            memset(vis, 0, sizeof vis);
            res = 0;
            dfs(1);
            if(res) printf("AYE\n");
            else printf("NAY\n");
        }
        printf(".\n");
    }
    return 0;
}