Cod sursa(job #2697515)

Utilizator Alex_tz307Lorintz Alexandru Alex_tz307 Data 18 ianuarie 2021 20:17:53
Problema Cerere Scor 100
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.3 kb
#include <bits/stdc++.h>
#define INF 0x3f3f3f3f

using namespace std;

ifstream fin("cerere.in");
ofstream fout("cerere.out");

int N, root;
vector<int> a, kth, S, dp;
vector<vector<int>> G;

void dfs1(int u, int parent) {
    if(a[u] != 0)
        kth[u] = S[(int)S.size() - a[u]];
    S.emplace_back(u);
    for(const int &v : G[u])
        if(v != parent)
            dfs1(v, u);
    S.pop_back();
}

void min_self(int &a, int b) {
    a = min(a, b);
}

void dfs2(int u, int parent) {
    if(a[u] != 0)
        min_self(dp[u], dp[kth[u]] + 1);
    for(const int &v : G[u])
        if(v != parent)
            dfs2(v, u);
}

int main() {
    fin >> N;
    a.resize(N + 1);
    for(int i = 1; i <= N; ++i)
        fin >> a[i];
    G.resize(N + 1);
    vector<bool> has_father(N + 1);
    for(int i = 1; i < N; ++i) {
        int u, v;
        fin >> u >> v;
        has_father[v] = true;
        G[u].emplace_back(v);
    }
    for(int i = 1; i <= N; ++i)
        if(!has_father[i]) {
            root = i;
            break;
        }
    kth.resize(N + 1);
    dfs1(root, 0);
    dp.resize(N + 1, INF);
    for(int i = 1; i <= N; ++i)
        if(a[i] == 0)
            dp[i] = 0;
    dfs2(root, 0);
    for(int i = 1; i <= N; ++i)
        fout << dp[i] << ' ';
}