Mancunian və rəngli ağac (Həll təhlili)

Təsəvvür edək ki, məsələnin sadələşdirilmiş forması mövcuddur. Gəlin ağacın bütün zirvələrinin bir rəngdə olduğunu fərz edək. Zəngin başlanğıcında boş olan bir yığın elan edək. Ağacın kökündən dərinlik üzrə axtarışa başlayaq. Zirvə $v$-ə daxil olarkən onu yığına əlavə edəcəyik, və zirvə $v$-ın emalı başa çatdıqda yığının üst elementini (bu zirvə $v$ olacaq) çıxaracağıq. Axtarış başa çatdıqdan sonra yığın yenə boş olacaq.

Sual: $v$ zirvəsinə daxil olarkən, yəni onu yığına yerləşdirməmişdən əvvəl yığının üzərində nə olacaq?

İndi isə əsas tapşırığın həllinə keçək. Rənglərin hər biri üçün (məsələn, yığın vektoru istifadə edərək) $c$ yığını yaradaq. Əvvəlcə bütün yığınlar boş olur. Ağacın kökündən — zirvə $1$-dən dərinlik üzrə axtarışa başlayaq. Rəngi $color$ ilə rənglənmiş $v$ zirvəsinin emalı aşağıdakı addımlardan ibarətdir:

• Əgər s[color] yığını boş deyilsə, onun üzərində rəngi $v$ zirvəsi ilə eyni olan ən yaxın ata zirvəsinin nömrəsi olacaq. Əgər yığın boşdursa, belə zirvə yoxdur və $v$ zirvəsi üçün cavab olaraq $-1$ çıxarılmalıdır.

• $v$ zirvəsini s[color] yığınına əlavə edirik.

• Zirvə $v$-ın bütün oğullarından rekursiv olaraq dərinlik üzrə axtarışa başlayırıq.

• Məlumatların işlənməsi bitdikdən sonra s[color] yığından $v$ zirvəsini çıxarırıq.

Ağacın dərinlik üzrə axtarışı zamanı, $v$ zirvəsində olan zaman yığınlarda kökdən $v$ zirvəsinədək olan yolda rənglər haqqında məlumatlar saxlanılır. Xüsusilə, s[color] yığını bu yolda $color$ rəngində olan zirvələrin nömrələrini saxlayır. Bu zirvələr dərinlik üzrə axtarış zamanı ziyarət edildiyindən ardıcıl olaraq yığına daxil edilir.

Nümunə

Dərinlik üzrə axtarış zirvə $5$-ə çatanda, c = 4 yığından ikisi boş olacaq (rəngləri $3$ və $4$ olanlar).

Rəng $1$ üçün yığın $1$ zirvəsini, rəng $2$ üçün yığın $3$ zirvəsini saxlayacaq. $5$ zirvəsi rəng $2$ olduğuna görə, ona ən yaxın eyni rəngdə olan ata, rəng $2$ yığınının üstündədir.

Beləliklə, rəng üzrə $5$ zirvəsinin ən yaxın atası $3$ zirvəsidir.

Alqoritmin reallaşdırılması

İşçi massivlər elan edək.

vector<int> col, res;
vector<vector<int> > g;
vector<stack<int> > s;

dfs funksiyası $v$ zirvəsindən başlayaraq dərinlik üzrə axtarışı yerinə yetirir.

void dfs(int v)
{

Zirvə $v$ rəngi $color$-dur.

  int color = col[v];

Rəngi $color$ olan $v$ zirvəsinin ən yaxın atası nömrəsini $res[v]$-də saxlayırıq.

  if(s[color].empty())
    res[v] = -1;
  else
    res[v] = s[color].top();

$v$ zirvəsini rəng $color$ yığına əlavə edirik.

  s[color].push(v);

$v$ zirvəsinin bütün $to$ oğullarını təkrar keçir və onlardan rekursiv olaraq dərinlik üzrə axtarışı başladır.

  for (int to : g[v])
    dfs(to);

Zirvə $v$-ın emalı başa çatdıqdan sonra (çıxış zamanı) onu rəng $color$ yığından silirik.

  s[color].pop();
}

Proqramın əsas hissəsi. Giriş məlumatlarını oxuyuruq. Ağac qururuq.

scanf("%d %d", &n, &c);
g.resize(n + 1);
for(i = 2; i <= n; i++)
{
  scanf("%d",&val);
  g[val].push_back(i);
}

Ağacın zirvələrinin rənglərini oxuyuruq.

col.resize(n + 1);
for(i = 1; i <= n; i++)
  scanf("%d",&col[i]);

Zirvə $1$-dən başlayan dərinlik axtarışına başlayırıq.

s.resize(c + 1);
res.resize(n + 1);
dfs(1);

Nəticəni çıxarırıq.

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

Python həyata keçirilməsi

Rekursiya üçün yığın həcmimizi artırırıq.

import sys
sys.setrecursionlimit(1000000)

dfs funksiyası $v$ zirvəsindən başlayaraq dərinlik üzrə axtarış yerinə yetirir.

def dfs(v):

Zirvə $v$ rəngi $color$-dur.

  color = col[v]

Rəngi $color$ olan $v$ zirvəsinin ən yaxın atası nömrəsini $res[v]$-də saxlayırıq.

  if not s[color]:
    res[v] = -1
  else:
    res[v] = s[color][-1]

$v$ zirvəsini rəng $color$ yığına əlavə edirik.

  s[color].append(v)

$v$ zirvəsinin bütün $to$ oğullarını təkrar keçir və onlardan rekursiv olaraq dərinlik üzrə axtarışın başladır.

  for to in g[v]:
    dfs(to)

Zirvə $v$-ın emalı başa çatdıqdan sonra (çıxış zamanı) onu rəng $color$ yığından silirik.

  s[color].pop()

Proqramın əsas hissəsi. Giriş məlumatlarını oxuyuruq. Ağac qururuq.

n, c = map(int, input().split())
lst = list(map(int, input().split()))
g = [[] for _ in range(n + 1)]
for i in range(2, n + 1):
  val = lst[i - 2]
  g[val].append(i)

Ağacın zirvələrinin rənglərini oxuyuruq.

col = [0] * (n + 1)
col[1:] = list(map(int, input().split()))

Zirvə $1$-dən başlayan dərinlik axtarışına başlayırıq.

s = [[] for _ in range(c + 1)]
res = [0] * (n + 1)
dfs(1)

Nəticəni çıxarırıq.

print(*res[1:])