Манкуніанець і кольорове дерево (Розбір)

Розглянемо спрощений варіант задачі. Нехай всі вершини дерева зафарбовані в один колір. Оголосимо стек, який спочатку пустий. Запустимо обхід у глибину з кореня дерева. При вході у вершину $v$ додамо її в стек, а при завершенні обробки вершини $v$ видалимо верхній елемент стека (ним якраз і буде вершина $v$). Після завершення обходу стек знову стане порожнім.

Запитання: що буде знаходитися на вершині стека в момент входу у вершину $v$, до того як ми помістимо $v$ в стек?

Тепер перейдемо до розв'язання вихідної задачі. Створимо $c$ стеків — по одному для кожного кольору (наприклад, можна використовувати вектор стеків). Спочатку всі стеки пусті. Запустимо обхід у глибину з кореня дерева — вершини $1$. Обробка вершини $v$, зафарбованої в колір $color$, включає наступні кроки:

• Якщо стек s[color] не пустий, то на його вершині знаходиться номер вершини, яка є найближчим предком вершини $v$ того ж кольору. Якщо стек пустий, то такої вершини не існує, і в якості відповіді для вершини $v$ слід вивести $-1$.

• Додаємо вершину $v$ в стек s[color].

• Рекурсивно запускаємо пошук у глибину від всіх синів вершини $v$.

• Після завершення обробки видаляємо вершину $v$ зі стека s[color].

Під час обходу у глибину, перебуваючи у вершині $v$, у стеках міститься інформація про кольори всіх вершин, розташованих на єдиному шляху від кореня дерева до вершини $v$. Зокрема, стек s[color] зберігає номери вершин на цьому шляху, зафарбованих у колір $color$. Ці вершини заносяться в стек у порядку їх відвідування при обході у глибину.

Приклад

Коли пошук у глибину дійде до вершини $5$, із c = 4 стеків два будуть порожніми (відповідні кольорам $3$ і $4$).

Стек номер $1$ (для кольору $1$) міститиме вершину $1$, стек номер $2$ (для кольору $2$) міститиме вершину $3$. Так як вершина $5$ має колір $2$, її найближчий предок з таким же кольором знаходиться на вершині стека номер $2$.

Таким чином, найближчим предком вершини $5$ за кольором є вершина $3$.

Реалізація алгоритму

Оголосимо робочі масиви.

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

Функція dfs виконує пошук у глибину, починаючи з вершини $v$.

void dfs(int v)
{

Вершина $v$ має колір $color$.

  int color = col[v];

Номер найближчого предка вершини $v$, що має колір $color$, зберігаємо в $res[v]$.

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

Заносимо вершину $v$ в стек номер $color$.

  s[color].push(v);

Перебираємо всіх синів $to$ вершини $v$ і рекурсивно запускаємо з них пошук у глибину.

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

Після завершення обробки вершини $v$ (тобто при виході) видаляємо її зі стека з номером $color$.

  s[color].pop();
}

Основна частина програми. Читаємо вхідні дані. Будуємо дерево.

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

Читаємо кольори вершин дерева.

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

Запускаємо пошук у глибину з вершини $1$.

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

Виводимо відповідь.

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

Python реалізація

Збільшуємо стек для рекурсії.

import sys
sys.setrecursionlimit(1000000)

Функція dfs виконує пошук у глибину, починаючи з вершини $v$.

def dfs(v):

Вершина $v$ має колір $color$.

  color = col[v]

Номер найближчого предка вершини $v$, що має колір $color$, зберігаємо в $res[v]$.

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

Заносимо вершину $v$ в стек номер $color$.

  s[color].append(v)

Перебираємо всіх синів $to$ вершини $v$ і рекурсивно запускаємо з них пошук у глибину.

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

Після завершення обробки вершини $v$ (тобто при виході) видаляємо її зі стека з номером $color$.

  s[color].pop()

Основна частина програми. Читаємо вхідні дані. Будуємо дерево.

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)

Читаємо кольори вершин дерева.

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

Запускаємо пошук у глибину з вершини $1$.

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

Виводимо відповідь.

print(*res[1:])