《P11249 [GESP202409 七级] 小杨寻宝》

题目背景

对应的选择、判断题：试题 - GESP 202409 C++ 七级 - 洛谷有题

题目描述

小杨有一棵包含 n 个节点的树，树上的一些节点放置有宝物。

小杨可以任意选择一个节点作为起点并在树上移动，但是小杨只能经过每条边至多一次，当小杨经过一条边后，这条边就会消失。小杨每经过一个放置有宝物的节点就会取得该宝物。

小杨想请你帮他判断自己能否成功取得所有宝物。

输入格式

本题单个测试点内有多组测试数据 。输入第一行包含一个正整数 t，代表测试用例组数。

接下来是 t 组测试用例。对于每组测试用例，一共 n+1 行。

第一行包含一个正整数 n，代表树的节点数。

第二行包含 n 个非负整数 a1,a2,...an，其中如果 ai=1，则节点 i 放置有宝物；若 ai=0，则节点 i 没有宝物。

之后 n−1 行，每行包含两个正整数 xi,yi，代表存在一条连接节点 xi 和 yi 的边。

输出格式

对于每组测试数据，如果小杨能成功取得所有宝物，输出 Yes，否则输出 No。

输入输出样例

输入 #1复制

复制代码

输出 #1复制

复制代码

Yes
No

说明/提示

样例 1 解释

对于第一组测试用例，小杨从节点 2 出发，按照 2−1−3−4 的顺序即可成功取得所有宝物。

数据规模与约定

子任务编号	数据点占比	t	n
1	20%	≤10	≤5
2	20%	≤10	≤103
3	60%	≤10	≤105

对全部的测试点，保证 1≤t≤10，1≤n≤105，0≤ai≤1，且保证树上至少有一个点放置有宝物。

代码实现：

cpp 复制代码

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100005;
const int LOG = 20;

vector<int> g[MAXN];
int a[MAXN];
int dep[MAXN];
int up[LOG][MAXN];

void dfs(int u,int f)
{
    up[0][u]=f;
    for(int i=1;i<LOG;i++) up[i][u]=up[i-1][up[i-1][u]];
    for(int v:g[u])
    {
        if(v==f) continue;
        dep[v]=dep[u]+1;
        dfs(v,u);
    }
}

int lca(int u,int v)
{
    if(dep[u]<dep[v]) swap(u,v);
    int dif = dep[u]-dep[v];
    for(int i=0;i<LOG;i++)
    {
        if(dif&(1<<i)) u=up[i][u];
    }
    if(u==v) return u;
    for(int i=LOG-1;i>=0;i--)
    {
        if(up[i][u] != up[i][v])
        {
            u=up[i][u];
            v=up[i][v];
        }
    }
    return up[0][u];
}

int get_dist(int u,int v)
{
    int anc = lca(u,v);
    return dep[u]+dep[v]-2*dep[anc];
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int t;
    cin>>t;
    while(t--)
    {
        int n;
        cin>>n;
        vector<int> treasure;
        for(int i=1;i<=n;i++)
        {
            g[i].clear();
            cin>>a[i];
            if(a[i]==1) treasure.push_back(i);
        }
        for(int i=1;i<=n-1;i++)
        {
            int u,v;
            cin>>u>>v;
            g[u].push_back(v);
            g[v].push_back(u);
        }
        int st0 = treasure[0];
        dep[st0]=0;
        dfs(st0,0);
        int s = st0;
        for(int x:treasure)
        {
            if(get_dist(st0,x) > get_dist(st0,s)) s=x;
        }
        dep[s]=0;
        dfs(s,0);
        int t_node = s;
        for(int x:treasure)
        {
            if(get_dist(s,x) > get_dist(s,t_node)) t_node = x;
        }
        int total_len = get_dist(s,t_node);
        bool ok = true;
        for(int x:treasure)
        {
            if(get_dist(s,x)+get_dist(x,t_node) != total_len)
            {
                ok = false;
                break;
            }
        }
        cout << (ok ? "Yes":"No") << '\n';
    }
    return 0;
}