Cod sursa(job #2560198)

Utilizator popashtefan10Popa Stefan popashtefan10 Data 27 februarie 2020 20:21:08
Problema Cerere Scor 30
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.32 kb
#include <iostream>
#include <cstdio>
#include <vector>
#define NMAX 100000

using namespace std;

int n;
int k[NMAX + 5], ans[NMAX + 5], log2[NMAX + 5], dad[NMAX + 5];
int rmq[17][NMAX + 5];
vector<int> v[NMAX + 5];

void init_log2() {
  log2[1] = 0;
  for(int i = 2; i <= n; i++)
    log2[i] = log2[i / 2] + 1;
}

void gen_rmq() {
  for(int i = 0; i <= log2[n]; i++)
    for(int j = 1; j <= n; j++)
      if(i == 0)
        rmq[i][j] = dad[j];
      else
        rmq[i][j] = rmq[i - 1][rmq[i - 1][j]];
}

int query(int x, int kx) {
  if(kx == 0)
    return x;

  int l = log2[kx];
  return query(rmq[l][x], kx - (1 << l));
}

void dfs(int x, int dx) {
  if(k[x] == 0)
    ans[x] = 0;
  else
    ans[x] = ans[query(x, k[x])] + 1;

  for(int y: v[x])
    if(y != dx)
      dfs(y, x);
}

int main() {
  freopen("cerere.in", "r", stdin);
  freopen("cerere.out", "w", stdout);
  int x, y;

  scanf("%d", &n);
  init_log2();
  for(int i = 1; i <= n; i++)
    scanf("%d", &k[i]);

  for(int i = 1; i < n; i++) {
    scanf("%d %d", &x, &y);
    dad[y] = x;
    v[x].push_back(y);
    v[y].push_back(x);
  }

  gen_rmq();

  for(int i = 1; i <= n; i++)
    if(dad[i] == 0) {
      dfs(i, 0);
      break;
    }

  for(int i = 1; i <= n; i++)
    printf("%d ", ans[i]);

  return 0;
}