Cod sursa(job #2781096)

Utilizator Tudor_PascaTudor Pasca Tudor_Pasca Data 8 octombrie 2021 14:25:58
Problema Cerere Scor 65
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.15 kb
#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

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

int n, lantLun, r;
int stramos[200100], ans[200100], lant[200100];
bool fr[200100];
vector<vector<int>> adj;

void dfs(int node, int father)
{
    lant[++lantLun] = node;

    if(stramos[node] != 0)
    {
        int indStramos = lantLun - stramos[node];
        int str = lant[indStramos];
        ans[node] = ans[str] + 1;
    }

    for(auto it: adj[node])
    {
        if(it == father)
            continue;

        dfs(it, node);
    }

    lantLun--;
}

int main()
{
    ios_base::sync_with_stdio(false);
    in >> n;

    for(int i = 1; i <= n; i++)
        in >> stramos[i];

    adj.resize(n + 5);
    for(int i = 1; i < n; i++)
    {
        int x, y;
        in >> x >> y;

        adj[x].push_back(y);
        adj[y].push_back(x);

        fr[y] = 1;
    }

    for(int i = 1; i <= n && r == 0; i++)
        if(fr[i] == 0)
            r = i;

    dfs(r, 0);

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

    out << '\n';

    return 0;
}