Cod sursa(job #2566733)

Utilizator luci.tosaTosa Lucian luci.tosa Data 3 martie 2020 00:54:24
Problema Cerere Scor 85
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.18 kb
#include <iostream>
#include <fstream>
#include <vector>
#define NMAX 100005
using namespace std;

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

vector <int> g[NMAX],modg[NMAX],roots;
int n,i,K[NMAX];
int stiva[NMAX],ans[NMAX];
bool root[NMAX];
int top,lvl;

void dfs(int nod) {
    if(!K[nod])
        roots.push_back(nod);
    else
        modg[stiva[top-K[nod]]].push_back(nod);
    for(int i=0; i<(int)g[nod].size(); i++) {
        int aux=g[nod][i];
        stiva[++top]=aux;
        dfs(aux);
        top--;
    }
}

void dfsmod(int nod) {
    ans[nod]=lvl;
    for(int i=0; i<(int)modg[nod].size(); i++) {
        int aux=modg[nod][i];
        lvl++;
        dfsmod(aux);
        lvl--;

    }
}
int main() {
    fin>>n;
    for(i=1; i<=n; i++) {
        fin>>K[i];
        root[i]=true;
    }
    for(i=1; i<n; i++) {
        int x,y;
        fin>>x>>y;
        g[x].push_back(y);
        root[y]=false;
    }
    for(i=1; i<=n; i++)
        if(root[i]==true) {
            stiva[++top]=i;
            break;
        }
    dfs(i);
    for(i=0; i<(int)roots.size(); i++) {
        lvl=0;
        dfsmod(roots[i]);
    }
    for(i=1; i<=n; i++)
        fout<<ans[i]<<" ";
    return 0;
}