3278392_tle.cc

北大大牛代码 1240道题的原代码 超级权威

#include <stdio.h>
#include <vector>
#include <string.h>
#include <stdlib.h>
#include <algorithm>
#define max 10005
using namespace std;

int n, f, r, m, ans, k;
int cnt[max], pos[max];
int visited[max], que[max], pa[max], checked[max], tque[max][2];
int flag;
typedef vector <int> arraylist;
vector <arraylist> sub;
arraylist pot, son;

struct node
{
	int w, id;
	int u, v;
}edge[max*2];

int cmp(const void *a,const void *b)
{
	struct node *aa = (struct node *)a;
	struct node *bb = (struct node *)b;

	return aa->u - bb->u;
}

arraylist getlist(int v,int MAX)
{
	int i;
	arraylist array;
	int front, rear;

	array.clear();
	flag++;
	front = 1;rear = 2;
	visited[v] = flag;
	tque[1][0] = v;
	tque[1][1] = 0;
	array.push_back(0);
	while(front < rear)
	{
		for(i = pos[tque[front][0]]; i < pos[tque[front][0]+1]; i++)
		{
			if(visited[edge[i].v]==flag)	continue;
			if(checked[edge[i].id]==1)	continue;
			visited[edge[i].v] = flag;
			int tmp = tque[front][1]+edge[i].w;
			if(tmp <= MAX)
			{
				tque[rear][0] = edge[i].v;
				tque[rear][1] = tmp;
				array.push_back(tmp);
				rear++;
			}
		}
		front++;
	}
	sort(array.begin(),array.end());
	return array;
}

void solve(int root)
{
	int i, j, tt;
	int min, minv, tmp;

	f = 1;
	r = 2;
	flag++;
	visited[root] = flag;
	que[1] = root;
	while(f < r)
	{
		for(i = pos[que[f]]; i < pos[que[f]+1]; i++)
		{
			if(visited[edge[i].v]==flag)	continue;
			if(checked[edge[i].id]==1)	continue;
			visited[edge[i].v] = flag;
			que[r] = edge[i].v;
			pa[edge[i].v] = que[f];
			r++;
		}
		f++;
	}
	if(r==2)
		return ;
	for(i = 1; i < r; i++)
	{
		cnt[que[i]] = 0;
	}
	for(i = r-1; i > 1; i--)
	{
		cnt[que[i]] ++;
		cnt[pa[que[i]]] += cnt[que[i]];
	}
	cnt[root]++;
	tmp = -1;
	for(j = pos[root]; j < pos[root+1]; j++)
	{
		if(cnt[edge[j].v] > tmp)
		{
			tmp = cnt[edge[j].v];
		}
	}
	min = tmp;minv = root;
	for(i = 2; i < r; i++)
	{
		tmp = -1;
		for(j = pos[que[i]]; j < pos[que[i]+1]; j++)
		{
			if(edge[j].v==pa[que[i]])
			{
				tt = n-cnt[que[i]];
			}
			else
			{
				tt = cnt[edge[j].v];
			}
			if(tt > tmp)
			{
				tmp = tt;
			}
		}
		if(tmp < min)
		{
			min = tmp;
			minv = que[i];
		}
	}
	sub.clear();
	pot.clear();
	son.clear();
	for(i = pos[minv]; i < pos[minv+1]; i++)
	{
		if(checked[edge[i].id]==1)
			continue;
		checked[edge[i].id] = 1;
		son.push_back(edge[i].v);
		if(edge[i].w > k)
			continue;
		arraylist t = getlist(edge[i].v,k-edge[i].w);
		sub.push_back(t);
		pot.push_back(edge[i].w);
		ans += t.size();
	}
	int size;
	size = sub.size();
	for(i = 0; i < size; i++)
	{
		for(j = i+1; j < size; j++)
		{
			int tmp = pot[i]+pot[j];
			if(tmp > k)
				continue;
			tmp = k-tmp;
			int i1, i2;
			for(i1 = 0; i1 < sub[i].size(); i1++)
			{
				int tt = sub[i][i1];
				if(tt > tmp)
					break;
				i2 = sub[j].size()-1;
				while(i2>-1&&sub[j][i2]+tt > tmp)
					i2--;
				if(i2==-1)
					break;
				else
					ans += i2+1;
			}
		}
	}
	size = son.size();
	for(i = 0; i < size; i++)
	{
		solve(son[i]);
	}
}

int main()
{
	int i, u, v, w;

	memset(visited,0,sizeof(visited));
	flag = 1;
	while(scanf("%d%d",&n,&k)==2)
	{
		if(n==0&&k==0)
		{
			break;
		}
		ans = 0;
		m = n-1;
		for(i = 1; i <= m; i++)
		{
			scanf("%d%d%d",&u,&v,&w);
			edge[i].id = edge[i+m].id = i;
			edge[i].u = edge[i+m].v = u;
			edge[i].v = edge[i+m].u = v;
			edge[i].w = edge[i+m].w = w;
		}
		m <<= 1;
		qsort(&edge[1],m,sizeof(edge[1]),cmp);
		for(i = m; i > 0; i--)
		{
			pos[edge[i].u] = i;
		}
		pos[n+1] = m+1;
		memset(checked,0,sizeof(checked));
		solve(1);
		printf("%d\n",ans);
	}
	return 0;
}