Cod sursa(job #2777711)

Utilizator lucamLuca Mazilescu lucam Data 23 septembrie 2021 22:36:03
Problema Cerere Scor 100
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 0.91 kb
#include <fstream>
#include <vector>

std::ifstream in("cerere.in");
std::ofstream out("cerere.out");

constexpr int N = 1e5 + 1;

std::vector<int> g[N];
int rg[N];
int v[N];
int dp[N];
int branch[N];

void dfs(int node, int level) {
    branch[level] = node;
    if (v[node]) {
        dp[node] = dp[branch[level - v[node]]] + 1;
    }
    for (int x : g[node]) {
        dfs(x, level + 1);
    }
}

int find_root(int node) {
    while (rg[node]) {
        node = rg[node];
    }
    return node;
}

int main() {
    int n;
    in >> n;
    int random_node;
    for (int i = 1; i <= n; ++i) {
        in >> v[i];
        if (!v[i]) {
            random_node = i;
        }
    }
    for (int i = 1; i < n; ++i) {
        int a, b;
        in >> a >> b;
        g[a].push_back(b);
        rg[b] = a;
    }
    dfs(find_root(random_node), 0);
    for (int i = 1; i <= n; ++i) {
        out << dp[i] << ' ';
    }
    out << '\n';
}