Cod sursa(job #2566731)

Utilizator luci.tosaTosa Lucian luci.tosa Data 3 martie 2020 00:43:39
Problema Cerere Scor 85
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.39 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];
int n,K[NMAX];
bool vis[NMAX];
int stiva[NMAX],ans[NMAX],parent[NMAX];
int top,lvl,root;

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

void dfsmod(int nod) {
    vis[nod]=0;
    ans[nod]=lvl;
    for(int i=0; i<(int)modg[nod].size(); i++) {
        int aux=modg[nod][i];
        if(vis[aux]) {
            lvl++;
            dfsmod(aux);
            lvl--;
        }
    }
}
int main() {
    fin>>n;
    for(int i=1; i<=n; i++) {
        fin>>K[i];
        parent[i]=-1;
    }
    for(int i=1; i<n; i++) {
        int x,y;
        fin>>x>>y;
        g[x].push_back(y);
        parent[y]=x;
    }
    for(int i=1; i<=n; i++)
        if(parent[i]==-1) {
            root=i;
            stiva[++top]=root;
        }
    dfs(root);
    for(int i=1;i<=n;i++) {
        if(parent[i]==-1) {
            lvl=0;
            dfsmod(i);
        }
    }
    for(int i=1;i<=n;i++)
        fout<<ans[i]<<" ";
    return 0;
}