Cod sursa(job #3302654)

Utilizator GliggyGligor Andrei Gliggy Data 9 iulie 2025 19:57:05
Problema Cerere Scor 85
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.31 kb
#include <bits/stdc++.h>
using namespace std;

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

const int N = 100005;
const int LOG = 20;

int n, k[N], par[N];
int dp[N][LOG];
int steps[N];

int get_kth_ancestor(int node, int jump) {
    for (int i = 0; i < LOG; ++i)
        if (jump & (1 << i))
            node = dp[node][i];
    return node;
}

int main() {
    ios::sync_with_stdio(false);
    fin.tie(nullptr); fout.tie(nullptr);

    fin >> n;
    for (int i = 1; i <= n; ++i) fin >> k[i];

    for (int i = 1; i < n; ++i) {
        int a, b;
        fin >> a >> b;
        par[b] = a;
    }

    // Build dp table
    for (int i = 1; i <= n; ++i)
        dp[i][0] = par[i];

    for (int j = 1; j < LOG; ++j)
        for (int i = 1; i <= n; ++i)
            dp[i][j] = dp[dp[i][j-1]][j-1];

    // Compute steps iteratively
    for (int i = 1; i <= n; ++i) {
        int curr = i;
        int cnt = 0;

        while (k[curr] != 0) {
            curr = get_kth_ancestor(curr, k[curr]);
            if (steps[curr]) {  // if already known, reuse
                cnt += steps[curr] + 1;
                break;
            }
            ++cnt;
        }
        steps[i] = cnt;
    }

    for (int i = 1; i <= n; ++i)
        fout << steps[i] << " ";

    return 0;
}