Cod sursa(job #2672986)

Utilizator StanCatalinStanCatalin StanCatalin Data 15 noiembrie 2020 16:33:35
Problema Cerere Scor 30
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 2.49 kb
#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

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

const int dim = 1e5 + 5;

int n,a[dim],stramosi[dim][18],dp[dim],lg[dim],radacina,am[dim];
vector<int> vec[dim];

int Get_Stramos(int nod,int k)
{
    ///cout << k << " ";
   /* if (k == 0) return nod;
    return Get_Stramos(stramosi[nod][lg[k]], k - (1<<lg[k]));*/

    int str = nod;
    while (k > 0)
    {
        str = stramosi[str][lg[k]];
        k -= (1<<lg[k]);
        ///cout << str << " " << k << "\n";
    }
    return str;
}

void DFS(int nod,int parent = 0)
{
    for (int k=1; k<18; k++)
    {
        stramosi[nod][k] = stramosi[stramosi[nod][k-1]][k-1];
        /// if (nod == 3 && k == 1) cout << stramosi[nod][k-1];
    }
    if (a[nod] > 0)
    {
        /// cout << nod << " " << a[nod] << " " << Get_Stramos(nod, a[nod]) << " " << dp[Get_Stramos(nod, a[nod])] << "\n";
        dp[nod] = dp[Get_Stramos(nod, a[nod])] + 1;
    }
    for (auto y:vec[nod])
    {
        if (y != parent)
        {
            DFS(y, nod);
        }
    }
}

void DFS_Rezolva(int nod,int parent = 0)
{
    if (a[nod] > 0)
    {
        /// cout << nod << " " << a[nod] << " " << Get_Stramos(nod, a[nod]) << " " << dp[Get_Stramos(nod, a[nod])] << "\n";
        dp[nod] = dp[Get_Stramos(nod, a[nod])] + 1;
    }
    for (auto y:vec[nod])
    {
        if (y != parent)
        {
            DFS_Rezolva(y, nod);
        }
    }
}

int main()
{
    in >> n;
    dp[1] = dp[0] = 0;
    for (int i=1; i<=n; i++)
    {
        in >> a[i];
    }
    stramosi[0][0] = 0;
    int x,y;
    for (int i=1; i<n; i++)
    {
        in >> x >> y;
       /// cout << x << " " << y << "\n";
        am[y] = 1;
        stramosi[y][0] = x;
        ///cout << y << " " << stramosi[y][0] << "\n";
        vec[x].push_back(y);
        vec[y].push_back(x);
    }
    int ok = 0;
    for (int i=1; i<=n && ok == 0; i++)
    {
        if (am[i] == 0)
        {
            radacina = i;
            ok = 1;
        }
    }
    for (int i=2; i<=n; i++)
    {
        lg[i] = 1 + lg[i/2];
    }
    DFS(radacina);
    ////for (int i=1; i<=n; i++) cout << stramosi[i][2] << " ";

   /// DFS_Rezolva(radacina);

    for (int i=1; i<=n; i++)
    {
        out << dp[i] << " ";
    }

  ///cout << Get_Stramos(6, 5);
    /// cout << stramosi[3][1];
    ///for (int i=1; i<=n; i++) cout << stramosi[i][1] << "\n";
    return 0;
}