【GMOJ6803】模拟tom

题目

题目链接：https://gmoj.net/senior/#main/show/6803
众所周知，Tom 猫对香肠非常感兴趣。
有一天，Tom 家里的女主人赏给了 Tom 一大堆香肠。这些香肠太多了，以至于 Tom 一顿吃不完，
于是它把这些香肠串成了一棵树，树的每个节点上都有一个香肠。
Tom 需要给这些香肠进行编号，其中有 a 个香肠需要编号为 1,2···a 中的不重复的编号，作为早餐
肠，剩下的 b 个香肠需要编号为 −1,−2··· − b 中的不重复的编号，作为晚餐肠。
Tom 每天会随机吃一顿饭，可能是早饭，也可能是晚饭。如果是吃早饭，Tom 会吃掉编号绝对值最
小的早餐肠，反之吃掉编号绝对值最小的晚餐肠。
如果一根香肠被吃掉了，那么与它相连的树上的边都会断掉，因此剩下的香肠可能会因此变成若干
棵树，即变得不再连通。这是 Tom 不希望发生的事。
请给这些香肠编号，使得无论 Tom 如何安排早饭和晚饭，整棵树一直都是连通的。

思路

首先如果早餐肠和晚餐肠不分别成连通块显然不行，只要将其中一方去完不取另一方就可以分成多个连通块。
否则按照拓扑序来编号一定有解。
而拆成两个连通块大小分别为 \(a\) 和 \(b\) 当且仅当存在一个子树大小为 \(a\) 或 \(b\)。将这个子树与其父亲的边断掉分别编号即可。
时间复杂度 \(O(n)\)。

代码

#include <bits/stdc++.h>
using namespace std;

const int N=100010;
int n,a,b,pos,last,tot=1,head[N],size[N],ans[N];
bool flag;

struct edge
{
	int next,to,from;
}e[N*2];

void add(int from,int to)
{
	e[++tot].to=to;
	e[tot].from=from;
	e[tot].next=head[from];
	head[from]=tot;
}

void dfs1(int x,int fa,int id)
{
	size[x]=1;
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa)
		{
			dfs1(v,x,i);
			size[x]+=size[v];
			if (pos) return;
		}
	}
	if (size[x]==a || size[x]==b) pos=id;
}

void dfs2(int x,int fa,int v)
{
	for (int i=head[x];~i;i=e[i].next)
	{
		int y=e[i].to;
		if (y!=fa) dfs2(y,x,v);
	}
	tot+=v; ans[x]=tot; last=tot;
}

int main()
{
	freopen("tom.in","r",stdin);
	freopen("tom.out","w",stdout);
	memset(head,-1,sizeof(head));
	scanf("%d%d%d",&n,&a,&b);
	for (int i=1,x,y;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
	}
	dfs1(1,0,0);
	if (!pos) return printf("-1"),0;
	tot=0; dfs2(e[pos].from,e[pos].to,1);
	tot=0; dfs2(e[pos].to,e[pos].from,-1);
	for (int i=1;i<=n;i++)
		if (last==-b) printf("%d\n",ans[i]);
			else printf("%d\n",-ans[i]);
	return 0;
}