Cod sursa(job #3159761)

Utilizator IvanAndreiIvan Andrei IvanAndrei Data 21 octombrie 2023 22:32:35
Problema Cerere Scor 100
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 2.15 kb
#include <bits/stdc++.h>

using namespace std;

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

int a[max_size], t[rmax][max_size], ans[max_size], uz[max_size];
pair <int, int> lvl[max_size];
vector <int> mc[max_size];

void dfs (int nod, int par, int niv)
{
    lvl[nod].first = nod;
    lvl[nod].second = niv;
    for (auto f : mc[nod])
    {
        if (f == par)
        {
            continue;
        }
        dfs(f, nod, niv + 1);
    }
}

bool cmp (pair <int, int> x, pair <int, int> y)
{
    return x.second < y.second;
}

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

void solve (int nod)
{
    if (ans[nod] != -1)
    {
        return;
    }
    if (a[nod] == 0)
    {
        ans[nod] = 0;
        return;
    }
    int anc = goup(nod, a[nod]);
    solve(anc);
    ans[nod] = 1 + ans[anc];
}

int main ()
{
    #ifdef LOCAL
      freopen("test.in", "r", stdin);
      freopen("test.out", "w", stdout);
    #else
      freopen("cerere.in", "r", stdin);
      freopen("cerere.out", "w", stdout);
    #endif // LOCAL
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    int n;
    cin >> n;
    for (int i = 1; i <= n; i++)
    {
        cin >> a[i];
        ans[i] = -1;
    }
    for (int i = 1; i < n; i++)
    {
        int x, y;
        cin >> x >> y;
        mc[x].push_back(y);
        t[0][y] = x;
        uz[y] = 1;
    }
    int rad = 0;
    for (int i = 1; i <= n; i++)
    {
        if (uz[i] == 0)
        {
            rad = i;
            break;
        }
    }
    dfs(rad, 0, 1);
    for (int e = 1; e < rmax; e++)
    {
        for (int i = 1; i <= n; i++)
        {
            t[e][i] = t[e - 1][t[e - 1][i]];
        }
    }
    sort(lvl + 1, lvl + n + 1, cmp);
    for (int i = 1; i <= n; i++)
    {
        if (ans[lvl[i].first] == -1)
        {
            solve(lvl[i].first);
        }
    }
    for (int i = 1; i <= n; i++)
    {
        cout << ans[i] << " ";
    }
    return 0;
}