Cod sursa(job #2310605)

Utilizator akaprosAna Kapros akapros Data 1 ianuarie 2019 17:45:48
Problema Cerere Scor 0
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.55 kb
#include <bits/stdc++.h>
#define maxN 100002
#define maxLog 18
using namespace std;

FILE *fin = freopen("cerere.in", "r", stdin);
FILE *fout = freopen("cerere.out", "w", stdout);

int n, lg2, r;
int kth[maxN];
vector < int > V[maxN];

bool root[maxN];
int anc[maxN][maxLog];

int dp[maxN];

void addEdge(int u, int v)
{
    -- u;
    -- v;
    V[u].push_back(v);
}
void init(int n)
{
    lg2 = 0;
    while ((1 << lg2) < n) ++ lg2;
    for (int i = 0; i < n; ++ i)
    {
        root[i] = true;
        scanf("%d", &kth[i]);
    }
}
int KthAnc(int nod, int k)
{
    for (int i = 0; (1 << i) <= k; ++ i)
        if (k & (1 << i))
            nod = anc[nod][i];
    return nod;
}
void dfs(int nod, bool ty)
{
    if (ty)
    {
        if (!kth[nod])
            dp[nod] = 0;
        else
            dp[nod] = dp[KthAnc(nod, kth[nod])] + 1;
    }
    for (int son : V[nod])
    {
        anc[son][0] = nod;
        dfs(son, ty);
    }

}
void compAnc()
{
    anc[r][0] = r;
    for (int j = 1; (1 << j) <= n; ++ j)
        for (int i = 0; i < n; ++ i)
            anc[i][j] = anc[anc[i][j - 1]][j - 1];
}
void writeAns()
{
    dfs(r, 1);
    for (int i = 0; i < n; ++ i)
        printf("%d ", dp[i]);
}

int main()
{
    scanf("%d", &n);
    init(n);
    for (int i = 1; i < n; ++ i)
    {
        int u, v;
        scanf("%d%d", &u, &v);
        root[v] = false;
        addEdge(u, v);
    }
    for (r = 0; r < n && !root[r]; ++ r);

    dfs(r, 0);
    compAnc();
    writeAns();
    return 0;
}