Cod sursa(job #2310586)

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

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

int n, lg2;
int kth[maxN];
int *anc[maxN];
int dp[maxN];
vector < int > V[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)
    {
        anc[i] = new int[lg2 + 1];
        for (int j = 0; j <= lg2; ++ j)
            anc[i][j] = 0;
        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)
        dp[nod] = dp[KthAnc(nod, kth[nod])] + (kth[nod] != 0);
    for (int son : V[nod])
    {
        anc[son][0] = nod;
        dfs(son, ty);
    }

}
void compAnc()
{
    for (int j = 1; j <= lg2; ++ j)
        for (int i = 1; i < n; ++ i)
            anc[i][j] = anc[anc[i][j - 1]][j - 1];
}

void writeAns()
{
    dfs(0, 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);
        addEdge(u, v);
    }
    dfs(0, 0);
    compAnc();
    writeAns();
    return 0;
}