Cod sursa(job #2764600)

Utilizator DragosC1Dragos DragosC1 Data 21 iulie 2021 19:59:40
Problema Cerere Scor 30
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.45 kb
#include <fstream>
#include <vector>
#include <iostream>
using namespace std;

int n;
int up[100001], fr[100001];
int dp[100001][17];
int radacina;
int rez[100001];

vector<int> a[100001];

void read() {
    int i, x, y;
    ifstream f("cerere.in");
    f >> n;
    for (i = 1; i <= n; i++)
        f >> up[i];
    for (i = 1; i < n; i++) {
        f >> x >> y;
        fr[y] = 1;
        a[x].push_back(y);
        a[y].push_back(x);
    }
    f.close();
}

void dfs(int x, int px) {
    int i, j;
    dp[x][0] = px;
    for (i = 1; i < 17; i++)
        dp[x][i] = dp[dp[x][i - 1]][i - 1];
    for (i = 0; i < a[x].size(); i++) {
        if (a[x][i] == px)
            continue;
        dfs(a[x][i], x);
    }
}

void solve() {
    int i;
    for (i = 1; i <= n; i++)
        if (fr[i] == 0)
            radacina = i;
    dfs(radacina, 0);
}

int calc(int x) {
    if (rez[x] != -1) 
        return rez[x];
    if (up[x] == 0)
        return 0;
    int node = x;
    for (int i = 0; i < 17; i++)
        if (up[x] & (1 << i)) 
            node = dp[node][i];
    if (up[node] == 0)
        return rez[x] = 1;
    else return rez[x] = 1 + calc(node);
}

void output() {
    int i;
    for (i = 1; i <= n; i++)
        rez[i] = -1;
    ofstream g("cerere.out");
    for (i = 1; i <= n; i++)
        g << calc(i) << ' ';
    g.close();
}

int main() {
    read();
    solve();
    output();
    return 0;
}