POJ 2114 Boatherds 划分树

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define MAX 10010
#define INF 0x3f3f3f3f
using namespace std;

int points,k;
int head[MAX],total;
int next[MAX << 1],aim[MAX << 1],length[MAX << 1];

int _total,_size,size[MAX];
int ans,root;
bool v[MAX];
int dis[MAX],p;

inline void Initialize();
inline void Add(int x,int y,int len);

void Work(int x);
void GetRoot(int x,int last);
inline int Count(int x,int last,int len);
void GetDis(int x,int last,int len);

int main()
{
while(scanf("%d",&points) != EOF && points) {
Initialize();
for(int y,z,i = 1;i <= points; ++i) {
while(scanf("%d",&y) != EOF && y) {
scanf("%d",&z);
Add(i,y,z),Add(y,i,z);
}
}
while(scanf("%d",&k) != EOF && k) {
ans = 0;
size[1] = points;
memset(v,false,sizeof(v));
Work(1);
if(ans)	puts("AYE");
else	puts("NAY");
}
puts(".");
}
return 0;
}

inline void Initialize()
{
total = 0;
memset(head,0,sizeof(head));
}

inline void Add(int x,int y,int len)
{
next[++total] = head[x];
aim[total] = y;
length[total] = len;
head[x] = total;
}

void Work(int x)
{
_size = INF,_total = size[x];
GetRoot(x,0);
x = root;
v[x] = true;
ans += Count(x,0,0);
for(int i = head[x];i;i = next[i]) {
if(v[aim[i]])	continue;
ans -= Count(aim[i],x,length[i]);
Work(aim[i]);
}
}

void GetRoot(int x,int last)
{
size[x] = 1;
int max_size = 0;
for(int i = head[x];i;i = next[i]) {
if(aim[i] == last || v[aim[i]])	continue;
GetRoot(aim[i],x);
size[x] += size[aim[i]];
max_size = max(max_size,size[aim[i]]);
}
max_size = max(max_size,_total - size[x]);
if(max_size < _size)
_size = max_size,root = x;
}

inline int Count(int x,int last,int len)
{
p = 0;
GetDis(x,last,len);
sort(dis + 1,dis + p + 1);
int l = 1,r = p,re = 0;
while(l < r) {
if(dis[l] + dis[r] > k)	--r;
else if(dis[l] + dis[r] < k)	++l;
else {
if(dis[r] == dis[l]) {
re += (r - l) * (r - l + 1) >> 1;
break;
}
int _l = l,_r = r;
while(dis[l] == dis[_l])	_l++;
while(dis[r] == dis[_r])	_r--;
re += (_l - l) * (r - _r);
l = _l,r = _r;
}
}
return re;
}

void GetDis(int x,int last,int len)
{
dis[++p] = len;
for(int i = head[x];i;i = next[i]) {
if(aim[i] == last || v[aim[i]])	continue;
GetDis(aim[i],x,len + length[i]);
}
}