【树形DP】AGC010C Cleaning

题目大意

Hellen和Shawn找到了一棵神奇的树.这棵树有n个节点,节点i有ai个石头.因为树上有石头是件奇怪的事情,两人决定将石头全部移除.Hellen觉得直接移除没有意思,所以对移除石头的操作进行了规定.
对于每次移除石头的操作,选择两个叶子节点(u,v)(u不能等于v),移除u到v路径上的每一个节点的一块石头(包括u,v).注意:如果这条路径上有一个节点没有石头,则不能进行操作.
此处的叶子节点为度数为1的节点.
Shawn想让你来告诉他, 能否通过上述操作将树上的石头移完.

解法

每个叶子结点向另一个叶子结点的连边，可拆成两条链，每一个关系在每一棵子树根统一处理。对于每一个结点x，记录一个d[x]表示该节点向上延伸的链有多少个，然后转移一下即可。

代码

#include<cstdio>
#define MAXN 100006
#include<algorithm>
#include<vector>
using namespace std;
vector<int> road[MAXN];
long long d[MAXN],a[MAXN];
int n,u,v,du[MAXN];
void DFS(int x,int father)
{
	if(du[x]==1)
	{
		d[x]=a[x];
		return;
	}
	long long sum=0,maxone=0;
	for(int i=0;i<(int)road[x].size();i++)
	{
		int nex=road[x][i];
		if(nex==father) continue;
		DFS(nex,x);
		sum+=d[nex];
		maxone=max(maxone,d[nex]);
	}
	if(a[x]>sum||a[x]*2<sum)
	{
		printf("NO");
		exit(0);
	}
	long long p=a[x]*2-sum;
	if(maxone-p>sum-maxone)
	{
		printf("NO");
		exit(0);
	}
	d[x]=p;
}
int main()
{
	//freopen("tree.in","r",stdin);
	//freopen("tree.out","w",stdout);
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
		scanf("%lld",&a[i]);
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&u,&v);
		road[u].push_back(v);
		road[v].push_back(u);
		du[u]++;
		du[v]++;
	}
	for(int i=1;i<=n;i++)
		if(du[i]!=1)
		{
			DFS(i,0);
			if(d[i]) printf("NO");
			else printf("YES");
			return 0;
		}
	if(a[1]==a[2]) printf("YES");
	else printf("NO");
}