Submission #73903788

---

Source Code Expand

#include <iostream>
#include <vector>
#include <unordered_set>
#include <stack>
using namespace std;
struct StackFrame {
    int node;
    int parent;
    bool visited;
    bool has_dup;
    bool inserted;
    StackFrame(int n, int p, bool v, bool hd, bool ins) 
        : node(n), parent(p), visited(v), has_dup(hd), inserted(ins) {}
};
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int N;
    cin >> N;
    vector<int> A(N + 1);
    for (int i = 1; i <= N; ++i) {
        cin >> A[i];
    }
    vector<vector<int>> adj(N + 1);
    for (int i = 0; i < N - 1; ++i) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    vector<string> ans(N + 1);
    unordered_set<int> path_nums;
    stack<StackFrame> st;
    st.emplace(1, -1, false, false, false);
    while (!st.empty()) {
        auto frame = st.top();
        st.pop();
        int u = frame.node;
        int parent = frame.parent;
        bool visited = frame.visited;
        bool has_dup = frame.has_dup;
        bool inserted = frame.inserted;
        if (!visited) {
            if (has_dup) {
                ans[u] = "Yes";
                st.emplace(u, parent, true, has_dup, false);
            } else {
                if (path_nums.count(A[u])) {
                    ans[u] = "Yes";
                    st.emplace(u, parent, true, true, false);
                } else {
                    ans[u] = "No";
                    path_nums.insert(A[u]);
                    st.emplace(u, parent, true, false, true);
                }
            }
            for (auto it = adj[u].rbegin(); it != adj[u].rend(); ++it) {
                int v = *it;
                if (v != parent) {
                    bool new_has_dup = (ans[u] == "Yes");
                    st.emplace(v, u, false, new_has_dup, false);
                }
            }
        } else {
            if (inserted) {
                path_nums.erase(A[u]);
            }
        }
    }
    for (int k = 1; k <= N; ++k) {
        cout << ans[k] << '\n';
    }
    return 0;
}

Submission Info

Judge Result