Знайдіть центроїд (Розбір)

Нехай маємо дерево з $n$ вершин.

Центроїдом дерева називається така вершина, при видаленні якої дерево розпадається на компоненти зв'язності, кількість вершин у кожній з яких не більше $n /2$ .

У задачі слід знайти всі центроїди дерева.

Розглянемо приклади дерев і їх центроїди (позначені зеленим кольором).

Запустимо пошук в глибину $df s (v)$ , який у $s u b [v]$ обчислить кількість вершин у піддереві з вершиною $v$ . Для наведених прикладів біля кожної вершини $v$ зазначимо відповідне їй значення $s u b [v]$ (пошук в глибину стартуємо у всіх прикладах з вершини $1$ ).

Якщо вершина $v$ має синів $t o_{1}, t o_{2}, ..., t o_{k}$ to_1, to_2, ..., to_k, то

s u b_{v} = 1 + s u b_{t o_{1}} + s u b_{t o_{2}} + ... + s u b_{t o_{k}}

sub_v = 1 + sub_{to_1} + sub_{to_2} + ... + sub_{to_k}

Вершина $v$ є центроїдом дерева тоді і тільки тоді, коли:

для кожного її сина $t o$ має місце $s u b [t o] \leq n /2$ sub[to] \le n / 2;
якщо з дерева видалити піддерево з коренем $v$ , то воно буде містити не більше $n /2$ вершин: $n - s u b [v] \leq n /2$ n - sub[v] \le n / 2;

Наприклад, у правому на рисунку дереві з $n = 5$ n = 5 вершинами центроїдом буде вершина $2$ , тому що

$s u b [4] = 2 \leq n /2$ sub[4] = 2 \le n / 2;
$s u b [3] = 1 \leq n /2$ sub[3] = 1 \le n / 2;
$n - s u b [2] = 5 - 4 = 1 \leq n /2$ n - sub[2] = 5 - 4 = 1 ≤ n / 2;

Приклад

Розглянемо дерево з прикладу. Вершина $6$ є центроїдом.

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

Граф будемо зберігати в списку суміжності $g$ .

vector<vector<int> > g;

Функція dfs повертає кількість вершин у піддереві з вершиною $v$ і зберігає це значення в $s u b [v]$ .

int dfs(int v, int p = -1)
{
  sub[v] = 1;
  for (int to : g[v])
    if (to != p) sub[v] += dfs(to, v);
  return sub[v];
}

C++

Функція centroid здійснює пошук в глибину, знаходить і заносить центроїди в масив $ce n t r$ .

void centroid(int v, int p = -1)
{

Встановимо $f l a g = 1$ flag = 1, якщо вершина $v$ є центроїдом.

  int flag = 1;

Перебираємо вершини $t o$ , суміжні з $v$ . Розглядаємо ребро $v \to t o$ v → to.

  for (int to : g[v])
    if (to != p)
    {

Якщо для сина $t o$ виконується $s u b [t o] > n /2$ sub[to] > n / 2, то $v$ не є центроїдом.

      if (sub[to] > n / 2) flag = 0;

Якщо для сина $t o$ виконується $s u b [t o] < n /2$ sub[to] < n / 2, то в піддереві з коренем $t o$ не міститься центроїдів. Є сенс продовжувати пошук у синові $t o$ , якщо тільки $s u b [t o] \geq n /2$ sub[to] \ge n / 2.

      if (sub[to] >= n / 2) centroid(to, v);
    }

Дерево без піддерева з коренем $v$ містить $n - s u b [v]$ n - sub[v] вершин. Якщо воно містить більше $n /2$ вершин, то $v$ не центроїд.

  if (n - sub[v] > n / 2) flag = 0;

Якщо для вершини $v$ виконуються всі умови центроїда, то заносимо її в масив $ce n t r$ .

  if (flag) centr.push_back(v);
}

Основна частина програми. Ініціалізуємо масиви.

scanf("%d", &n);
g.resize(n + 1);
sub.resize(n + 1);

Читаємо вхідний граф.

for (i = 0; i < n - 1; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}

Запускаємо пошук в глибину, шукаємо центроїди дерева.

dfs(1);
centroid(1);

Виводимо один із центроїдів.

printf("%d\n", centr[0]);

Реалізація алгоритму — друге рішення

Запустимо пошук в глибину $df s (v)$ , який для кожної вершини $v$ обчислить наступні величини:

$s u b [v]$ містить кількість вершин у піддереві з вершиною $v$ ;
$f [v]$ містить максимальний розмір піддерева серед усіх піддерев вершини $v$ , включаючи піддерево, що містить батьківську вершину $v$ .

Тоді центроїдом буде вершина $v$ з найменшим значенням $f [v]$ . Таких вершин може бути або одна або дві в дереві.

Вершина $v = 6$ v = 6 з'єднана з $4$ піддеревами, розміри яких відповідно рівні $1, 2, 3$ і $5$ . Значення $f [6] = 5$ f[6] = 5 містить максимальний розмір піддерева.

Водночас найменшим значенням у масиві $f$ є $f [6] = 5$ f[6] = 5. Отже, вершина $6$ є центроїдом.

Розглянемо розстановку міток для дерева з двома центроїдами.

Дві вершини мають найменше значення в масиві $f$ : $f [1] = 3$ f[1] = 3 і $f [2] = 3$ f[2] = 3.

void dfs(int v, int p = -1)
{
  sub[v] = 1;
  f[v] = 0;
  for (int to : g[v])
    if (to != p)
    {
      dfs(to, v);
      sub[v] += sub[to];
      f[v] = max(f[v], sub[to]);
    }
  f[v] = max(f[v], n - sub[v]);
  if ((root == 0) || (f[v] < f[root])) root = v;
}

C++

Основна частина програми. Читаємо вхідні дані. Ініціалізуємо масиви.

scanf("%d", &n);
g.resize(n + 1);
f.resize(n + 1);
sub.resize(n + 1);
for (i = 0; i < n - 1; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}

C++

Запускаємо пошук в глибину, виводимо один із центроїдів дерева $r oo t$ .

root = 0;
dfs(1);
printf("%d\n", root);

Python реалізація

Збільшимо розмір стека.

import sys
sys.setrecursionlimit(300000)

Функція dfs повертає кількість вершин у піддереві з вершиною $v$ і зберігає це значення в $s u b [v]$ .

def dfs(v, p = -1):
  sub[v] = 1
  for to in g[v]:
    if to != p:
      sub[v] += dfs(to, v)
  return sub[v]

Функція centroid здійснює пошук в глибину, знаходить і заносить центроїди в масив $ce n t r$ .

def find_centroid(v, p = -1):

Встановимо $f l a g = T r u e$ flag = True, якщо вершина $v$ є центроїдом.

  flag = True

Перебираємо вершини $t o$ , суміжні з $v$ . Розглядаємо ребро $v \to t o$ v → to.

  for to in g[v]:
    if to == p: continue

Якщо для сина $t o$ виконується $s u b [t o] > n /2$ sub[to] > n / 2, то $v$ не є центроїдом.

    if sub[to] > n // 2:
      flag = False

    if sub[to] >= n // 2:
      find_centroid(to, v)

  if n - sub[v] > n // 2:
      flag = False

Якщо для вершини $v$ виконуються всі умови центроїда, то заносимо її в масив $ce n t r$ .

  if flag: centr.append(v)

Основна частина програми. Читаємо вхідні дані. Ініціалізуємо масиви.

n = int(input())
g = [[] for _ in range(n + 1)]
sub = [0] * (n + 1)
centr = []

Читаємо вхідний граф.

for i in range(n - 1):
  a, b = map(int, input().split())
  g[a].append(b)
  g[b].append(a)

Запускаємо пошук в глибину, шукаємо центроїди дерева.

dfs(1)
find_centroid(1)

Виводимо один із центроїдів.

print(centr[0])

Python реалізація — друге рішення

Збільшимо розмір стека.

import sys
sys.setrecursionlimit(200000)

Запустимо пошук в глибину $df s (v)$ , який для кожної вершини $v$ обчислить наступні величини:

$s u b [v]$ містить кількість вершин у піддереві з вершиною $v$ ;
$f [v]$ містить максимальний розмір піддерева серед усіх піддерев вершини $v$ , включаючи піддерево, що містить батьківську вершину $v$ .

Тоді центроїдом буде вершина $v$ з найменшим значенням $f [v]$ . Таких вершин може бути або одна або дві в дереві.

def dfs(v, p=-1):
  global root
  sub[v] = 1
  f[v] = 0
  for to in g[v]:
    if to != p:
      dfs(to, v)
      sub[v] += sub[to]
      f[v] = max(f[v], sub[to])
  f[v] = max(f[v], n - sub[v])
  if root == 0 or f[v] < f[root]:
    root = v

Python

Основна частина програми. Читаємо вхідні дані. Ініціалізуємо масиви.

n = int(input())
g = [[] for _ in range(n + 1)]
f = [0] * (n + 1)
sub = [0] * (n + 1)

for _ in range(n - 1):
  a, b = map(int, input().split())
  g[a].append(b)
  g[b].append(a)

Python

Запускаємо пошук в глибину, виводимо один із центроїдів дерева $r oo t$ .

root = 0
dfs(1)
print(root)