Cod sursa(job #2710153)

Utilizator PatrickCplusplusPatrick Kristian Ondreovici PatrickCplusplus Data 22 februarie 2021 00:05:43
Problema Cerere Scor 30
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.16 kb
#include <bits/stdc++.h>

using namespace std;

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

const int nmax = 100005;
int n, k[nmax], radacina[nmax], dp[nmax][18], cnt[nmax];
vector <int> G[nmax];

int getParent(int nod, int p){
    int i = 0;
    while (p){
        if (p & 1) nod = dp[nod][i];
        p >>= 1;
        ++i;
    }
    return nod;
}

void dfs(int nod){
    if (k[nod] != 0){
        cnt[nod] = 1 + cnt[getParent(nod, k[nod])];
    }
    for (auto it : G[nod]){
        dfs(it);
    }
}

int main(){
    fin >> n;
    for (int i = 1; i <= n; ++i){
        fin >> k[i];
    }
    for (int i = 1; i < n; ++i){
        int x, y;
        fin >> x >> y;
        radacina[y] = 1;
        G[x].push_back(y);
        dp[y][0] = x;
    }
    int root;
    for (int i = 1; i <= n; ++i){
        if (radacina[i] == 0){
            root = i;
            break;
        }
    }
    for (int i = 1; i <= 16; ++i){
        for (int j = 1; j <= n; ++j){
            dp[j][i] = dp[dp[j][i - 1]][i - 1];
        }
    }
    dfs(root);
    for (int i = 1; i <= n; ++i){
        fout << cnt[i] << " ";
    }
    return 0;
}