3983. Манкунианец и цветное дерево

После напряженной недели на работе, жители Манчестера и Ливерпуля решили отправиться в поход на выходные. Прогуливаясь по лесу, они наткнулись на уникальное дерево, состоящее из n вершин. Вершины дерева пронумерованы числами от 1 до n.

Каждой вершине дерева присвоен один из c возможных цветов. Чтобы побороть скуку, они решили проверить свои логические навыки. Корнем дерева является вершина 1. Для каждой вершины они решили найти ближайшего предка, цвет которого совпадает с цветом этой вершины.

Вход. Первая строка содержит два целых числа n и c (1 ≤ n, c ≤ 10⁵) – количество вершин в дереве и количество возможных цветов.

Вторая строка содержит n – 1 целых чисел, где i-ое число указывает на отца (i + 1) - ой вершины.

Третья строка содержит n целых чисел, определяющих цвета вершин. Значения цветов лежат в диапазоне от 1 до c включительно.

Выход.  Выведите в одной строке n чисел, где i-ое число – это вершина, являющаяся ближайшим предком i-ой вершины с таким же цветом. Если такого предка для вершины не существует, выведите -1.

Пример входа	Пример выхода
5 4 1 1 3 3 1 4 2 1 2	-1 -1 -1 1 3

РЕШЕНИЕ

графы – поиск в глубину

Анализ алгоритма

Рассмотрим более простой вариант задачи. Пусть все вершины дерева покрашены в один цвет. Объявим стек, который изначально пуст. Запустим поиск в глубину из корня дерева. При входе в вершину v кладем значение v в стек, а при завершении обработки вершины v удаляем верхушку стека (это как раз будет вершина v). Когда поиск в глубину завершится, стек окажется пустым.

Вопрос: что будет находиться на вершине стека при входе в вершину v, до того как мы положим v в стек?

Теперь перейдем к решению нашей задачи. Создадим c стеков – по одному для каждого цвета (например, используем вектор стеков). Изначально все стеки пусты. Запустим поиск в глубину из корня – вершины 1. Обработка вершины v цвета color состоит из следующих шагов:

· Если стек s[color] не пуст, то на его вершине находится номер вершины, являющейся ближайшим предком вершины v с таким же цветом, как у v. Если стек пуст, то требуемой вершины не существует, и ответом для вершины v будет -1.

· Добавляем вершину v в стек s[color].

· Рекурсивно запускаем поиск в глубину со всех сыновей вершины v.

· Удаляем вершину v из стека s[color].

Когда при поиске в глубину мы находимся в вершине v, в стеках хранится информация о цветах всех вершин, расположенных на единственном пути от корня до v. То есть стек s[color] содержит номера вершин на пути от корня до v, которые имеют цвет color. При этом вершины в стек заносятся в порядке их посещения во время поиска в глубину.

Пример

Когда поиск в глубину дойдет до вершины 5, из c = 4 стеков два будут пустыми (соответствующие цветам 3 и 4). Стек номер 1 (для первого цвета) будет содержать вершину 1, стек номер 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. Извлекаем 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)

def dfs(v):

color = col[v]

if not s[color]:

res[v] = -1

else:

res[v] = s[color][-1]

s[color].append(v)

for to in g[v]:

dfs(to)

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()))

s = [[] for _ in range(c + 1)]

res = [0] * (n + 1)

dfs(1)

print(*res[1:])