Cod sursa(job #2316131)

Utilizator TheNextGenerationAyy LMAO TheNextGeneration Data 11 ianuarie 2019 10:56:17
Problema Cerere Scor 55
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.6 kb
#include <bits/stdc++.h>

using namespace std;

const int N = 1e5+1;
const int L = 20;

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

int dp[L][N],val[N],rez[N],lvl[N];
bool a[N];
vector<int> v[N];

void dfs(int x)
{
    for (auto it: v[x])
        if (!lvl[it])
        {
            lvl[it] = lvl[x]+1;
            dp[0][it] = x;
            dfs(it);
        }
}

int query(int x, int p)
{
    int L = 1, dif = lvl[x]-p;
    while ((1<<L)<=lvl[x])
        L++;
    for (int i = L; i>=0; i--)
        if (lvl[x]-(1<<i)>=dif)
            x = dp[i][x];
    return x;
}

int main()
{
    int n;
    in >> n;
    for (int i = 1; i<=n; i++)
        in >> val[i];
    for (int i = 1; i<n; i++)
    {
        int x,y;
        in >> x >> y;
        v[x].push_back(y);
        a[y] = 1;
    }
    int root = 0;
    for (int i = 1; i<=n; i++)
        if (!a[i])
            root = i;
    dfs(root);
    for (int i = 1; (1<<i)<=n; i++)
        for (int j = 1; j<=n; j++)
            dp[i][j] = dp[i-1][dp[i-1][j]];
    queue<int> q;
    lvl[root] = 1;
    q.push(root);
    memset(a,0,sizeof(a));
    a[root] = 1;
    while (!q.empty())
    {
        int now = q.front();
        q.pop();
        for (auto it: v[now])
            if (!a[it])
            {
                a[it] = 1;
                q.push(it);
                if (val[it])
                {
                    int above = query(it,val[it]);
                    rez[it] = rez[above]+1;
                }
            }
    }
    for (int i = 1; i<=n; i++)
        out << rez[i] << " ";
}