Cod sursa(job #1864735)

Utilizator MithrilBratu Andrei Mithril Data 31 ianuarie 2017 22:47:32
Problema Cerere Scor 0
Compilator cpp Status done
Runda Arhiva de probleme Marime 1.19 kb
#include <bits/stdc++.h>

using namespace std;
ifstream fin("cerere.in");
ofstream fout("cerere.out");
const int NMAX = 100005;
int n,x,y;
int kcost[NMAX],lev[NMAX];
vector<int> G[NMAX];
int dp[20][NMAX];
int memo[NMAX];
void DFS(int,int);
int getParent(int,int);

void DFS(int n,int l)
{
    lev[n]=l;
    for(auto w:G[n])
        DFS(w,l+1);
}

int getParent(int nod,int l)
{
    if(memo[nod]!=-1) return memo[nod];

    for(int p=0 ; (1<<p)<=l ; p++)
        if(l&(1<<p))
            nod=dp[p][nod];
    if(kcost[nod])
    {
        int r = getParent(nod,kcost[nod]);
        memo[nod] = r;
        return r;
    }
    return nod;
}

int main()
{
    fin>>n;
    for(int i=1;i<=n;i+=1) fin>>kcost[i];
    for(int i=1;i<=n;i+=1) memo[i]=-1;
    for(int i=1;i<n;i+=1)
    {
        fin>>x>>y;
        dp[0][y]=x;
        G[x].push_back(y);
    }

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

    DFS(1,1);
    for(int i=1;i<=n;i+=1)
    {
        int aux = lev[getParent(i,kcost[i])];
        aux = abs(lev[i]-aux);
        aux = (aux<=1) ? aux:aux-1;
        fout << aux << ' ';
    }
    return 0;
}