Cod sursa(job #2651554)

Utilizator felix24mihaiPrava Felix Mihai felix24mihai Data 22 septembrie 2020 22:14:56
Problema Cerere Scor 35
Compilator cpp-32 Status done
Runda Arhiva de probleme Marime 1.51 kb
#include <iostream>
#include <fstream>
#include <vector>
#define NMAX 100003
using namespace std;
ifstream f("cerere.in");
ofstream g("cerere.out");
int n, ancestorsLvl[NMAX], father[NMAX], answer[NMAX], head, visited[NMAX];
vector <int> nodes[NMAX];
void read(){
    int a,b;
    f >> n;
    long long s = n * (n+1) / 2;
    for (int i = 0; i < n; i++){
        f >> ancestorsLvl[i+1];
        if (ancestorsLvl[i+1] == 0)
            answer[i+1] = 0;
        else
            answer[i+1] = -1;
    }
    for (int i = 1; i < n; i++){
        f >> a >> b;
        s -= b;
        father[b] = a;
        nodes[a].push_back(b);
        nodes[b].push_back(a);
    }
    head = s;
}
int ancesLvl(int node, int lvl){
    for (int i = 0; i < lvl; i++)
        node = father[node];
    return node;
}
int resolve(int node, int nr){
    if (ancestorsLvl[node] == 0)
        return nr;
    else{
        int nextNode = ancesLvl(node,ancestorsLvl[node]);
        if (answer[nextNode] != -1){
            answer[node] = answer[nextNode] + 1;
            return answer[nextNode] + 1;
        }
        return resolve(nextNode, nr+1);
    }
}
void DFS(int node){
    if (answer[node] == -1)
        answer[node] = resolve(node,0);
    visited[node] = 1;
    for (int i = 0; i < nodes[node].size(); i++)
        if (visited[nodes[node][i]] == 0)
            DFS(nodes[node][i]);
}
int main()
{
    read();
    DFS(head);
    for (int i = 1; i <= n; i++)
        g << answer[i] << " ";
    return 0;
}