【AtCoder2304】Cleaning

题意

一棵树有n个节点,节点i有A[i]个石头。将石头全部移除。
对于每次移除石头的操作,选择两个叶子节点(u,v)(u不能等于v),移除u到v路径上的每一个节点的一块石头(包括u,v).注意:如果这条路径上有一个节点没有石头,则不能进行操作.
（此处的叶子节点为度数为1的节点）
能否通过上述操作将树上的石头移完.

题解

寻找一个度数大于1的结点为根，dfs下去，每次返回，这棵子树上还需多少次操作。
当前结点为u，处理完u的所有儿子的子树，这些子树还剩下一些操作，需要通过u来完成。
剩下需要的操作数两两配对，即组成了若干条路径，最后使得子树剩余操作数与当前结点u的石子数相同，剩余的石子需要更上层的路径消除。

设儿子子树剩余sum次操作，两两配对时，每次从儿子子树中取出两次操作配对，当前结点消掉一个石子，则总操作次数为sum-A[u]。
sum<A[u]无解（操作次数为负）
A[u]<sum-A[u]无解（当前结点石子不够）
取mx为儿子子树剩余操作数的最大值
sum-mx<sum-A[u]无解（无法两两配对）

对于根节点，需使得剩余操作数为0，则必须满足A[u]==sum/2

如果n==2，特判（找不到根。。。）

代码

#include<cstdio> #include<cstdlib> #include<vector> #include<algorithm> using namespace std; const int MAXN=100005; void NoSolution() {
      puts("NO"); fclose(stdin); fclose(stdout); exit(0); } int n,A[MAXN]; vector<int> adj[MAXN]; int dfs(int u,int fa=0) {
      if(adj[u].size()==1U) return A[u]; long long sum=0; int tmp=0,mx=0; for(int i=0;i<(int)adj[u].size();i++) {
      int v=adj[u][i]; if(v==fa) continue; tmp=dfs(v,u); mx=max(mx,tmp); sum+=tmp; } if(fa==0) {
      if(sum%2==1||mx>sum/2||A[u]!=sum/2) NoSolution(); return 0; } long long delta=sum-A[u]; if(A[u]<delta||sum<A[u]||sum-mx<delta) NoSolution(); return A[u]-delta; } int main() {
      scanf("%d",&n); for(int i=1;i<=n;i++) scanf("%d",&A[i]); if(n==2) {
      if(A[1]==A[2]) puts("YES"); else puts("NO"); return 0; } for(int i=1;i<n;i++) {
      int a,b; scanf("%d%d",&a,&b); adj[a].push_back(b); adj[b].push_back(a); } for(int i=1;i<=n;i++) if(adj[i].size()>1U) {
      dfs(i); break; } puts("YES"); return 0; }