Cod sursa(job #2777798)

Utilizator CaptnBananaPetcu Tudor CaptnBanana Data 24 septembrie 2021 18:48:29
Problema Cerere Scor 40
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.39 kb
#include <bits/stdc++.h>

using namespace std;

ifstream f("cerere.in");
ofstream g("cerere.out");

const int N = 1e5 + 1;
int n, k[N], x, y, ans[N], tata[N], h[N], noduri[N];
vector<int> c[N];

int stramos(int nod, int dist){
    if(dist == 0)
        return ans[nod] + 1;

    return stramos(tata[nod], dist - 1);
}

void dfs1(int nod, int distDeLaRadacina){
    h[nod] = distDeLaRadacina;
    for(auto y: c[nod])
        dfs1(y, distDeLaRadacina + 1);
}

void dfs2(int nod){
    if(ans[nod] != -1 && ans[nod] != 0)
        return;

    ans[nod] = stramos(nod, k[nod]);
    for(auto y: c[nod]){
        dfs2(y);
    }
}

bool cmp(int a, int b){
    return h[a] < h[b];
}

int main(){
    f >> n;
    for(int i = 1; i <= n; i++){
        ans[i] = -1;
        f >> k[i];
        if(k[i] == 0)
            ans[i] = 0;
    }

    for(int i = 1; i < n; i++){
        f >> x >> y;
        tata[y] = x;
        c[x].push_back(y);
    }

    f.close();
    for(int i = 1; i <= n; i++){
        if(tata[i] == 0){
            dfs1(i, 0);
            break;
        }
    }

    for(int i = 1; i <= n; i++)
        noduri[i] = i;

    sort(noduri + 1, noduri + n + 1, cmp);
    for(int i = 1; i <= n; i++){
        if(ans[noduri[i]] == 0)
            dfs2(noduri[i]);
    }

    for(int i = 1; i <= n; i++)
        g << ans[i] - 1 << ' ';

    g.close();
}