Cod sursa(job #2778104)

Utilizator pielevladutPiele Vladut Stefan pielevladut Data 28 septembrie 2021 16:05:09
Problema Cerere Scor 0
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.6 kb
#include <bits/stdc++.h>

using namespace std;

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

int n;
int k[100100];
int stramosi[100100][20];
///stramosi[i][j] -> al 2 ^ j lea stramos al lui i
vector<int> v[100100];

void dfs(int nod, int tata)
{
    if(tata != 0)
    {
        stramosi[nod][0] = tata;
    }
    for(auto it : v[nod])
    {
        if(it != tata)
        {
            dfs(it, nod);
        }
    }
}

int kth(int nod, int k)
{
    int ans = nod;
    for(int i = 0; (1 << i) <= k; i ++)
    {
        /// inainte puneam if(k & (1 << i) != 0)
        /// si nu efectueaza bine operatiile pe biti
        if((k & (1 << i)) != 0)
        {
            ans = stramosi[ans][i];
        }
    }
    return ans;
}

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;
        v[x].push_back(y);
        v[y].push_back(x);
    }
    dfs(1, 0);
    for(int i = 1; i <= n; i ++)
    {
        for(int j = 1; (1 << j) <= n; j ++)
        {
            stramosi[i][j] = stramosi[stramosi[i][j-1]][j-1];
        }
    }
    for(int i = 1; i <= n; i ++)
    {
        if(k[i] == 0)
        {
            fout << 0 << ' ';
        }
        else
        {
            int parinte = kth(i, k[i]);
            int ans = 1;
            while(k[parinte] != 0)
            {
                ans++;
                parinte = kth(parinte, k[parinte]);
            }
            fout << ans << ' ';
        }

    }
    return 0;
}