Cod sursa(job #2916807)

Utilizator IvanAndreiIvan Andrei IvanAndrei Data 1 august 2022 17:20:25
Problema Cerere Scor 40
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.63 kb
#include <fstream>
#include <vector>
#include <algorithm>

using namespace std;

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

const int max_size = 1e5 + 1, max_r = 18;

int t[max_r][max_size], a[max_size], ans[max_size];
bool tata[max_size];
/// cel care ramane 0 e radacina => pt dfs si nivelare
vector <int> edges[max_size];
vector <pair <int, int>> lvl;

void dfs (int nod, int nivel)
{
    lvl.push_back({nivel, nod});
    for (auto f : edges[nod])
    {
        dfs(f, nivel + 1);
    }
}

int anc (int x, int ord)
{
    int e = 0;
    while (ord > 0)
    {
        if (ord % 2 == 1)
        {
            x = t[e][x];
        }
        e++;
        ord /= 2;
    }
    return x;
}

int main ()
{
    int n;
    in >> n;
    for (int i = 1; i <= n; i++)
    {
        in >> a[i];
    }
    for (int i = 1; i < n; i++)
    {
        int x, y;
        in >> x >> y;
        edges[x].push_back(y);
        t[0][y] = x;
        tata[y] = true;
    }
    for (int e = 1; e < max_r; e++)
    {
        for (int i = 1; i <= n; i++)
        {
            t[e][i] = t[e - 1][t[e - 1][i]];
        }
    }
    for (int i = 1; i <= n; i++)
    {
        if (!tata[i])
        {
            dfs(i, 1);
            break;
        }
    }
    sort(lvl.begin(), lvl.end());
    for (int i = 1; i < n; i++)
    {
        if (a[lvl[i].second] > 0)
        {
            ans[lvl[i].second] = 1 + ans[anc(lvl[i].second, a[lvl[i].second])];
        }
    }
    for (int i = 1; i <= n; i++)
    {
        out << ans[i] << " ";
    }
    in.close();
    out.close();
    return 0;
}