题解：P11249 [GESP202409 七级] 小杨寻宝

题目显然等价于问所有宝箱是否在一条链上。

稍微转化一下题意，即我们现在要找到一条链，使得这条链上有宝物的节点数量尽可能多。

想到这里我们发现这个和树的直径比较相似，那么我们直接大胆将深度定义为从根到这个节点上有宝箱节点的数量，然后做一遍树的直径。最后判断直径长度是否等于所有有宝箱节点的数量即可。

做完以后我们发现这道题过了，其实证明也很简单。我们考虑把价值转到边上，即对于所有有宝箱的节点，让它所连的所有边长度都加上一。因为如果一条路径经过这条边，那么它一定也经过这个点，所以这样给边赋权后跑树的直径是对的，即上述思路正确。

Code

#include<bits/stdc++.h>
//#include <ext/pb_ds/assoc_container.hpp>
//#include <ext/pb_ds/tree_policy.hpp>
#define int long long							
using namespace std;
//using namespace  __gnu_pbds;
//tree<int,null_type,less<int>,rb_tree_tag,tree_order_statistics_node_update> tr;//从小到大
//int findnum(int k){auto it=tr.find_by_order(k-1);return ((it!=tr.end())?(*it):1e9+7);}//查元素
//int findrank(int x){return tr.order_of_key(x)+1;}//查排名
inline int read()
{
	int w=1,s=0;char ch=getchar();
	while(!isdigit(ch)){if(ch=='-')w=-1;ch=getchar();}
	while(isdigit(ch)){s=s*10+(ch-'0');ch=getchar();}
	return w*s;
}
const int maxn=3e5+10;
const int mod=998244353;
const int inf=1e9+7;
int n,a[maxn],dep[maxn];
vector<int> G[maxn];
void dfs(int x,int fa)
{
	dep[x]=dep[fa]+(a[x]==1);
	for(auto y : G[x])
	{
		if(y==fa)continue;
		dfs(y,x);
	}
}
void Main()
{
	n=read();
	for(int i=1;i<=n;i++)dep[i]=0,G[i].clear();
	int sum=0;
	for(int i=1;i<=n;i++)a[i]=read(),sum+=a[i];
	for(int i=1;i<n;i++)
	{
		int x=read(),y=read();
		G[x].push_back(y);
		G[y].push_back(x);
	}
	dfs(1,0);
	int ma=-1,id=0;
	for(int i=1;i<=n;i++)
	{
		if(dep[i]>ma)
		{
			ma=dep[i];
			id=i;
		}
	}
	for(int i=1;i<=n;i++)dep[i]=0;
	dfs(id,0);
	ma=-1;
	for(int i=1;i<=n;i++)
	{
		ma=max(ma,dep[i]);
	}
	puts(ma==sum?"Yes":"No");
}
signed main()
{
#ifdef Lydic
    freopen(".in","r",stdin);
    freopen(".out","w",stdout);
#endif
	int T=read();
	while(T--)Main();
    return 0;
}