Разбор задачи «Уровни старого дерева»

1. Идея

Нужно для каждой вершины дерева найти расстояние от заданной вершины r.

Так как все рёбра дерева имеют одинаковый вес 1, задача сводится к поиску кратчайших расстояний от одной вершины во невзвешенном графе. Для этого лучше всего подходит обход в ширину — BFS.

BFS запускается из вершины r и посещает вершины слоями:

• сначала саму вершину r на расстоянии 0,
• потом все вершины на расстоянии 1,
• потом все вершины на расстоянии 2,
• и так далее.

Именно это и требуется по условию: уровень вершины — число рёбер на пути от r до неё.

2. Наблюдения

1. Граф является деревом, значит:

   • в нём n вершин и n - 1 ребро;
   • между любыми двумя вершинами существует ровно один простой путь.

2. Поэтому расстояние от r до любой вершины определено однозначно.

3. Если мы впервые пришли в вершину to из вершины v, то её уровень равен dist[v] + 1.

4. Чтобы не заходить в вершины повторно, удобно хранить массив dist:

   • dist[x] = -1, если вершина ещё не посещена;
   • иначе там уже записан её уровень.

5. Так как дерево может содержать до 100000 вершин, нужно решение за O(n). BFS как раз работает за линейное время.

3. Алгоритм

1. Считать n и r.
2. Построить список смежности g, где g[v] содержит всех соседей вершины v.
3. Создать массив dist длины n + 1 и заполнить его значениями -1.
4. Положить dist[r] = 0, потому что расстояние от корня до самого себя равно 0.
5. Создать очередь и добавить в неё вершину r.
6. Пока очередь не пуста:
   • взять вершину v из начала очереди;
   • пройти по всем её соседям to;
   • если dist[to] == -1, значит, вершина ещё не посещена:
     • записать dist[to] = dist[v] + 1;
     • добавить to в очередь.
7. Вывести значения dist[1], dist[2], ..., dist[n].

4. Почему это работает

BFS обходит граф по слоям расстояний.

Сначала в очереди только вершина r, и для неё dist[r] = 0. Это верно.

Предположим, мы достали из очереди вершину v, для которой уже правильно найдено расстояние от r. Тогда для любого непосещённого соседа to:

• существует путь до to длины dist[v] + 1 через вершину v;
• более короткого пути быть не может, потому что BFS посещает вершины в порядке неубывания расстояния.

Значит, когда мы впервые записываем dist[to] = dist[v] + 1, это и есть правильное расстояние.

Так как дерево связно, BFS из вершины r рано или поздно посетит все вершины. Следовательно, для каждой вершины будет найден её правильный уровень.

5. Сложность

Пусть n — число вершин.

• Построение списка смежности: O(n)
• BFS: O(n), потому что каждая вершина обрабатывается один раз, и каждое ребро рассматривается не более двух раз

Итоговая сложность: O(n)

Память:

• список смежности: O(n)
• массив расстояний: O(n)
• очередь: O(n)

Итого по памяти: O(n)

6. Код на C++17

#include <iostream>
#include <vector>
#include <queue>

using namespace std;

int main() {
    int n, r;
    cin >> n >> r;

    vector<vector<int>> g(n + 1);
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    vector<int> dist(n + 1, -1);
    queue<int> q;
    dist[r] = 0;
    q.push(r);

    while (!q.empty()) {
        int v = q.front();
        q.pop();

        for (int to : g[v]) {
            if (dist[to] == -1) {
                dist[to] = dist[v] + 1;
                q.push(to);
            }
        }
    }

    for (int i = 1; i <= n; i++) {
        if (i > 1) {
            cout << ' ';
        }
        cout << dist[i];
    }
    cout << '\n';

    return 0;
}

7. Код на Python 3

from collections import deque

n, r = map(int, input().split())

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

dist = [-1] * (n + 1)
dist[r] = 0

q = deque([r])

while q:
    v = q.popleft()
    for to in g[v]:
        if dist[to] == -1:
            dist[to] = dist[v] + 1
            q.append(to)

print(*dist[1:])