Cod sursa(job #2907018)

Utilizator amcbnCiobanu Andrei Mihai amcbn Data 28 mai 2022 14:29:57
Problema Cerere Scor 75
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.28 kb
/// [A][M][C][B][N] ///
#include <bits/stdc++.h>
using namespace std;
const int mod = 1e9 + 7, inf = 0x3f3f3f3f;
const char sp = ' ', nl = '\n';
ifstream fin("cerere.in");
ofstream fout("cerere.out");

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    int n;
    fin >> n;
    vector<int> v(n + 1), dp(n + 1), mark(n + 1, 1);
    vector<vector<int>> g(n + 1);
    for (int i = 1; i <= n; ++i) {
        fin >> v[i];
    }
    for (int i = 1; i < n; ++i) {
        int x, y;
        fin >> x >> y;
        g[x].push_back(y);
        mark[y] = 0;
    }
    int root = 0;
    for (int i = 1; i <= n; ++i) {
        if (mark[i]) {
            root = i;
        }
    }
    stack<int> st;
    st.push(root);
    vector<int> check(n + 1), path;
    while (st.size()) {
        int node = st.top();
        if (!check[node]) {
            if (v[node]) {
                dp[node] = dp[path[path.size() - v[node]]] + 1;
            }
            path.push_back(node);
            for (auto& son : g[node]) {
                st.push(son);
            }
            check[node] = 1;
        }
        else {
            st.pop();
            path.pop_back();
        }
    }
    for (int i = 1; i <= n; ++i) {
        fout << dp[i] << sp;
    }
}