A1810 [GESP-7.2409.T2] 小杨寻宝

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int t;
    cin >> t;
    while (t--) {
        int n;
        cin >> n;
        vector<int> a(n + 1);
        int total = 0;
        for (int i = 1; i <= n; i++) {
            cin >> a[i];
            total += a[i];
        }

        vector<vector<int>> g(n + 1);
        for (int i = 1; i < n; i++) {
            int x, y;
            cin >> x >> y;
            g[x].push_back(y);
            g[y].push_back(x);
        }

        // 非递归遍历，求出每个节点的父亲以及遍历序
        vector<int> parent(n + 1, 0), order;
        order.reserve(n);
        order.push_back(1);
        parent[1] = -1;
        for (size_t i = 0; i < order.size(); i++) {
            int u = order[i];
            for (int v : g[u]) {
                if (v == parent[u]) continue;
                parent[v] = u;
                order.push_back(v);
            }
        }

        // cnt[u]：u 的子树里宝物的个数
        vector<int> cnt(n + 1, 0);
        for (int i = (int)order.size() - 1; i >= 0; i--) {
            int u = order[i];
            cnt[u] += a[u];
            if (parent[u] > 0) cnt[parent[u]] += cnt[u];
        }

        // 树上不走回头路的走法一定是一条简单路径，
        // 所以要求：每个点出发的各个方向上，最多只有两个方向有宝物
        bool ok = true;
        for (int u = 1; u <= n && ok; u++) {
            int dirs = 0;
            for (int v : g[u]) {
                if (parent[v] == u && cnt[v] > 0) dirs++;
            }
            if (total - cnt[u] > 0) dirs++;
            if (dirs > 2) ok = false;
        }

        cout << (ok ? "Yes" : "No") << "\n";
    }
    return 0;
}