Запустим пошук в глибину з вершини . При вході в вершину кольору color[v] збільшимо значення used[color[v]] на . Значення used[color[v]] містить, скільки разів колір color[v] зустрівся на шляху від вершини до вершини (включно). Якщо колір color[v] на цьому шляху зустрівся лише один раз, значить вершина є хорошою, і ми додаємо її в результатуючу множину.
При виході з вершини значення used[color[v]] слід зменшити на .
Приклад
Граф з першого прикладу виглядає наступним чином:
Вершина не є хорошою, так як на шляху 1 - 2 - 5 вершини та мають однаковий колір.
Вершина є хорошою, так як на шляху 1 - 2 - 3 - 6 кольори всіх проміжних вершин відрізняються від кольору вершини .
Реалізація алгоритму
Вхідний граф зберігаємо в списку суміжності . Оголосимо робочі масиви.
vector<int> used, color;
vector<vector<int>> g;
set<int> st;Функція dfs реалізує пошук в глибину. Змінна є предком вершини .
void dfs(int v, int par)
{Вершина має колір color[v]. Позначаємо, що на шляху від вершини зустрілася вершина цього кольору.
used[color[v]]++;Значення used[color[v]] показує, скільки разів колір color[v] зустрівся на шляху від вершини до вершини (включно). Якщо цей колір зустрівся на шляху рівно один раз, то вершина вважається хорошою, і ми додаємо її в результатуючу множину .
if (used[color[v]] == 1) st.insert(v);Перебираємо всі ребра (v, to), виходячі з вершини . Якщо вершина не є предком v~(to \ne par), запускаємо з неї пошук в глибину. При цьому предком стає .
for (int to : g[v])
if (to != par) dfs(to, v);При виході з вершини зменшуємо значення used[color[v]] на .
used[color[v]]--;
}Основна частина програми. Читаємо кількість вершин та масив кольорів.
scanf("%d", &n);
color.resize(n + 1);
for (i = 1; i <= n; i++)
scanf("%d", &color[i]);Читаємо граф.
used.resize(100001);
g.resize(n + 1);
for (i = 1; i < n; i++)
{
scanf("%d %d", &x, &y);
g[x].push_back(y);
g[y].push_back(x);
}Запускаємо пошук в глибину з вершини .
dfs(1, 1);Виводимо хороші вершини. Вони містяться в множині .
for (int val : st)
printf("%d\n", val);