Pythonist

Двоичные деревья вечны. По крайней мере, так думают технические менеджеры, занимающиеся наймом разработчиков. И когда на техническом собеседовании вас просят решить задачу, касающуюся двоичных деревьев, первое, что интервьюер захочет знать, — в ширину или в глубину?

В чем разница?

Когда говорят о поиске в ширину (Breadth First Search, BFS) и глубину (Depth First Search, DFS), имеется в виду порядок обхода узлов двоичного дерева. При обходе в глубину вы сначала опускаетесь к низу дерева, а потом идете в сторону, а при обходе в ширину — наоборот, начинаете с корня и спускаетесь сначала к его узлам-потомкам, обходите их, потом спускаетесь к потомкам потомков, обходите их, и так далее.

Если взять в качестве примера двоичное дерево на этом рисунке, при BFS-подходе порядок обхода узлов будет следующим: 1, 2, 3, 4, 5.

В случае с DFS возможны разные варианты последовательности посещения узлов. Все зависит от того, будет это прямой, обратный или центрированный обход. Например, прямой обход выдаст 1, 2, 4, 5, 3. Мы разбирали эти три типа обхода в статье «Обход двоичного дерева на Python».

Реализация поиска в ширину на Python

Давайте посмотрим, как сделать BFS на Python. Начнем с определения нашего класса TreeNode. В нем должны быть только значение и левый и правый указатели.

class TreeNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

1. Стратегия

Вот как мы будем обходить узлы:

1. Найдем высоту дерева от корня до самого дальнего листа.
2. В цикле переберем все уровни (заканчивая верхушкой).
3. Обойдем каждый уровень, выводя все узлы.

2. Поиск высоты дерева

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

def height(node):

Да, вы угадали: метод будет рекурсивным, как и большинство методов для обхода двоичных деревьев. Давайте подумаем, каким будет базовый случай. Самый простой вариант — когда у нас нулевой корень. В этом случае высота — 0. Вам может показаться, что простейший случай — узел без потомков, но тогда нам придется проверять, есть ли потомки, а это уже увеличивает сложность.

def height(node):
    if not node:
        return 0

Что дальше? Нам нужно обойти левую и правую половину дерева. Мы будем вызывать метод height() для этих половин и сохранять в две переменные: l_height и r_height.

def height(node):
    if not node:
        return 0
    l_height = height(node.left) 
    r_height = height(node.right) 

Какую высоту мы вернем: левую или правую? Конечно, большую! Поэтому мы берем max() обоих значений и возвращаем его. Не забудьте добавить 1 для текущего узла.

def height(node):
    if not node:
        return 0
    l_height = height(node.left) 
    r_height = height(node.right) 

    return max(l_height, r_height) + 1

3. Перебор уровней в цикле

Теперь, когда у нас есть высота, мы можем приступить к написанию метода для обхода дерева в ширину. Единственный аргумент, который этот метод будет принимать, — корневой узел.

def breadth_first(root):

Затем нам нужно получить нашу высоту.

def breadth_first(root):
    h = height(root)

Наконец, мы перебираем в цикле высоту дерева. На каждой итерации мы будем вызывать вспомогательную функцию — print_level(), которая принимает узел root и уровень.

def breadth_first(root):
    h = height(root)
    for i in range(h):
        print_level(root, i)

4. Обход дерева

Мы будем обходить каждый уровень отдельно, выводя все узлы на нем. Тут нам пригодится наша вспомогательная функция. Она принимает два аргумента: корень и текущий уровень.

def print_level(root, level):

В этом методе номер уровня определяется индексом i цикла for. Чтобы обойти все дерево, мы будем рекурсивно двигаться по уровням, уменьшая i, пока не достигнем уровня 0. Когда достигнем, это будет означать, что теперь можно выводить узлы на экран.

Наш базовый случай — когда корень — null. Этот метод только выводит дерево на экран и ничего не возвращает, поэтому в базовом случае будет просто return.

def print_level(root, level):
    if not root:
        return

Далее, если наш уровень — 0 (то есть интересующий нас уровень), нам нужно вывести значение этого узла.

def print_level(root, level):
    if not root:
        return
    if level == 0:
        print(root.value)

Наконец, если номер уровня больше нуля, это означает, что нам нужно продолжить обход дерева. Мы вызываем наш рекурсивный метод для левой и правой половин дерева. Не забудьте отнять единицу от level в обоих вызовах функции.

def print_level(root, level):
    if not root:
        return
    if level == 0:
        print(root.value)
    elif level > 0:
        print_level(root.left, level - 1)
        print_level(root.right, level - 1)

Вот и все! Нам пришлось создать три разных метода, но в конечном итоге мы можем обойти дерево в ширину.

5. Тестирование

Для начала давайте создадим дерево, используя наш класс TreeNode.

tree = TreeNode(1)

Теперь нужно заполнить оставшуюся часть дерева. При помощи следующих строк кода мы приведем дерево к тому виду, что изображен на схеме.

tree.left = TreeNode(2)
tree.right = TreeNode(3)
tree.left.left = TreeNode(4)
tree.left.right = TreeNode(5) 

Наконец, запускаем наш метод.

breadth_first(tree)

Ожидаемый результат: 1, 2, 3, 4, 5. Мы успешно обошли дерево!