poj 2114 Boatherds（树分治）

点分治，和poj 1741 Tree基本一样，就是统计数量那里不一样。

#include <stdio.h>
#include <string.h>
#include <algorithm>
using namespace std;

const int MAXN = 10010;
const int MAXM = 20020;
struct Edge
{
    int to,w,next;
} edge[MAXM];
int head[MAXN];
int son[MAXN];
int depth[MAXN];
int dr;
int n,k,tot,tmd;
bool vis[MAXN];

void init()
{
    tot = 0;
    memset(son,0,sizeof(son));
    memset(head,-1,sizeof(head));
}

void addedge(int u, int v, int w)
{
    edge[tot].to = v;
    edge[tot].w = w;
    edge[tot].next = head[u];
    head[u] = tot++;
}

void getRoot(int N, int u, int fa, int &root)
{
    son[u] = 1;
    int ans = 0;
    for(int i = head[u]; i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        if(v == fa || vis[v]) continue;
        getRoot(N,v,u,root);
        son[u] += son[v];
        ans = max(ans,son[v]);
    }
    ans = max(ans,N-son[u]);
    if(ans < tmd)
    {
        tmd = ans;
        root = u;
    }
}

void getDepth(int u, int fa, int dep)
{
    depth[dr++] = dep;
    for(int i = head[u]; i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        int w = edge[i].w;
        if(v == fa || vis[v]) continue;
        getDepth(v,u,dep+w);
    }
}

int calc(int u, int dep)
{
    dr = 0;
    int ret = 0;
    getDepth(u,-1,dep);
    sort(depth,depth+dr);
    int i = 0, j = dr-1;
    //统计depth[i]+depth[j] <= k的对数
    while(i < j)
    {
        while(depth[i]+depth[j] > k && i < j) --j;
        ret += (j-i);
        ++i;
    }
    i = 0, j = dr-1;
    //减去depth[i]+depth[j] < k的对数
    while(i < j)
    {
        while(depth[i]+depth[j] >= k && i < j) --j;
        ret -= (j-i);
        ++i;
    }
    return ret;
}

int solve(int u, int N)
{
    int root,ret = 0;
    tmd = MAXN;
    getRoot(N,u,-1,root);
    vis[root] = true;
    ret += calc(root,0);
    for(int i = head[root]; i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        if(vis[v]) continue;
        int w = edge[i].w;
        ret -= calc(v,w);
        ret += solve(v,son[v]);
    }
    return ret;
}

int main()
{
    int a,b,w;
    while(scanf("%d",&n) && n)
    {
        init();
        for(int i = 1; i <= n; ++i)
        {
            while(scanf("%d",&a) && a)
            {
                scanf("%d",&b);
                addedge(i,a,b);
                addedge(a,i,b);
            }
        }
        while(scanf("%d",&w) && w)
        {
            k = w;
            memset(vis,0,sizeof(vis));
            if(solve(1,n)) puts("AYE");
            else puts("NAY");
        }
        puts(".");
    }
    return 0;
}