Дерево#2923

Дерево (Розбір)

Виконуємо обхід дерева в глибину, починаючи з кореня. Для кожної вершини $i$ будемо зберігати множину $s_{i}$ s_i, у якій накопичуються кольори всіх вершин, що входять у піддерево з коренем в $i$ .

Якщо $j$ — син вершини $i$ в процесі обходу, то множина $s_{j}$ s_j повинна бути включена в $s_{i}$ s_i.

Кількість різних кольорів у піддереві з коренем у вершині $i$ дорівнює потужності (розміру) множини $s_{i}$ s_i.

Приклад

Ліворуч для кожної вершини вказано її колір, праворуч — множина кольорів, що містяться в піддереві з коренем у цій вершині.

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

У масиві $co l or [i]$ color[i] зберігається колір вершини $i$ . У множині $s [i]$ будемо накопичувати кольори, що зустрічаються в піддереві з коренем у вершині $i$ .

Орієнтоване дерево представлено у вигляді списку суміжності графа $g$ . У масиві $r es [i]$ зберігається кількість різних кольорів у піддереві з коренем у вершині $i$ .

#define MAX 1000010
int color[MAX], res[MAX];
set<int> s[MAX];
vector<vector<int> > g;

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

void dfs(int v)
{

Спочатку у множину $s [v]$ додаємо колір самої вершини $v$ .

  s[v].insert(color[v]);

Перебираємо ребра дерева $(v, t o)$ (v, to).

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

Для кожного ребра $(v, t o)$ (v, to) додаємо множину $s [t o]$ у $s [v]$ . Якщо розмір множини $s [v]$ менший за розмір множини $s [t o]$ , то міняємо їх місцями. Далі вміст меншої множини $s [t o]$ додаємо у множину $s [v]$ .

    if (s[v].size() < s[to].size()) s[v].swap(s[to]);
    s[v].insert(s[to].begin(), s[to].end());

Очищаємо множину $s [t o]$ — вона нам більше не знадобиться.

    s[to].clear();
  }

Після цього у $r es [v]$ записуємо кількість різних кольорів у піддереві, тобто розмір множини $s [v]$ .

  res[v] = s[v].size();
}

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

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

C++

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

dfs(0);

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

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