首页 > 代码库 > POJ 2114 Boatherds 树的分治

POJ 2114 Boatherds 树的分治

题目大意:给出一棵树,问有没有两点之间的距离是k的。多组数据


思路:和IOI2011的Race一样,比那个简单。读入太恶心了,我是上网上抄的别人的主函数。


CODE:


#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]);
	}
}


POJ 2114 Boatherds 树的分治