GESP2024年9月认证C++七级( 第三部分编程题（1）小杨寻宝）

参考程序：

#include <bits/stdc++.h>
using namespace std;
const int N = 1e5+10;
vector<int> g[N];  // 图的邻接表
int col[N], dep[N], has[N];

// 深度优先遍历，计算每个节点的深度
void dfs(int x, int fa) {
    dep[x] = dep[fa] + 1;  // 计算当前节点的深度
    for (auto i : g[x]) {  // 遍历与x相邻的节点
        if (i != fa) {
            dfs(i, x);  // 递归调用
        }
    }
}

// 判断从某个节点开始是否能够收集所有宝物
bool dfs2(int x, int fa) {
    for (auto i : g[x]) {
        if (i != fa) {
            auto res = dfs2(i, x);  // 递归判断子树
            if (res == false) return false;  // 如果有子树不能满足条件，返回false
            if (has[i] || col[i]) has[x]++;  // 如果子树或者该节点放有宝物，则当前节点记录已收集的宝物数量
        }
    }
    if (has[x] > 1) return false;  // 如果某节点收集到两个宝物，说明不能收集所有宝物
    return true;
}

int main() {
    int t;
    cin >> t;  // 输入测试用例数
    while (t--) {
        int n;
        cin >> n;  // 输入节点数
        for (int i = 1; i <= n; i++) {
            dep[i] = 0;
            g[i].clear();
            cin >> col[i];  // 输入每个节点是否有宝物
        }
        
        // 输入树的边
        for (int i = 1; i < n; i++) {
            int u, v;
            cin >> u >> v;
            g[u].push_back(v);
            g[v].push_back(u);
        }
        
        // 从节点1开始进行深度优先搜索，计算每个节点的深度
        dfs(1, 0);
        
        // 找到放置宝物的最远节点
        int mx = 0, pos = 0;
        for (int i = 1; i <= n; i++) {
            has[i] = 0;  // 初始化has数组
            if (col[i] && dep[i] > mx) {
                mx = dep[i];  // 更新最远的放置宝物的节点
                pos = i;  // 记录最远的宝物节点
            }
        }
        
        // 从最远的宝物节点出发，进行第二次深度优先搜索
        bool res = dfs2(pos, 0);
        if (res) {
            cout << "Yes\n";  // 如果能收集所有宝物，输出Yes
        } else {
            cout << "No\n";  // 否则输出No
        }
    }
}