Ультимативный гайд по структурам данных и алгоритмам на Python. Часть 5: деревья, BST и heapq

После поиска и сортировки в линейных коллекциях (Часть 4), пора перейти к нелинейным структурам. Файловые системы, DOM, оргструктуры — всё это иерархии, для моделирования которых в программировании служат деревья (trees). Понимание деревьев — ключ к эффективной работе с иерархическими данными и многим продвинутым алгоритмам.

Сегодня мы разберем основы деревьев, сфокусируемся на бинарных деревьях, научимся их обходить (DFS/BFS), изучим бинарные деревья поиска (BST), коснёмся куч (heaps) и модуля `heapq`.

Введение в деревья

Бинарные деревья: не более двух потомков ✌️

Реализация узла

class TreeNode:
    """Представляет узел бинарного дерева."""
    def __init__(self, key):
        self.key = key     # Данные узла
        self.left = None   # Ссылка на левого потомка
        self.right = None  # Ссылка на правого потомка

    def __repr__(self):
        """Строковое представление узла (показывает ключ)."""
        return f"Node({self.key})"

Создание дерева вручную

      *      <-- Корень
     / \
    +   3    <-- Узел '+' и лист '3'
   / \
  1   2      <-- Листья '1' и '2'

# 1. Узлы-листья (операнды)
node1, node2, node3 = TreeNode(1), TreeNode(2), TreeNode(3)
# 2. Узлы-операторы
node_plus, node_multiply = TreeNode('+'), TreeNode('*')

# 3. Строим связи
node_plus.left, node_plus.right = node1, node2
node_multiply.left, node_multiply.right = node_plus, node3

# 4. Корень дерева
root_expression_tree = node_multiply
print(f"Корень дерева: {root_expression_tree}")
# Пример доступа: root_expression_tree.left -> Node(+), root_expression_tree.left.left -> Node(1)

Как обойти дерево?

Обход в глубину (DFS)

def inorder(node):
    """In-order (LNR) обход."""
    return inorder(node.left) + [node.key] + inorder(node.right) if node else []

def preorder(node):
    """Pre-order (NLR) обход."""
    return [node.key] + preorder(node.left) + preorder(node.right) if node else []

def postorder(node):
    """Post-order (LRN) обход."""
    return postorder(node.left) + postorder(node.right) + [node.key] if node else []

# --- Тестовое дерево ---
#       F
#      / \
#     B   G
#    / \   \
#   A   D   I
#      / \   /
#     C   E H
root = TreeNode('F')
root.left = TreeNode('B'); root.right = TreeNode('G')
root.left.left = TreeNode('A'); root.left.right = TreeNode('D')
root.left.right.left = TreeNode('C'); root.left.right.right = TreeNode('E')
root.right.right = TreeNode('I'); root.right.right.left = TreeNode('H')

print(f"In-order   (LNR): {inorder(root)}")
# ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']
print(f"Pre-order  (NLR): {preorder(root)}")
# ['F', 'B', 'A', 'D', 'C', 'E', 'G', 'I', 'H']
print(f"Post-order (LRN): {postorder(root)}")
# ['A', 'C', 'E', 'D', 'B', 'H', 'I', 'G', 'F']

Обход в ширину (BFS)

from collections import deque

def level_order(root):
    """BFS (Level-Order) обход."""
    if not root: return []
    result = []
    nodes_queue = deque([root])
    while nodes_queue:
        current_node = nodes_queue.popleft()
        result.append(current_node.key)
        if current_node.left:
            nodes_queue.append(current_node.left)
        if current_node.right:
            nodes_queue.append(current_node.right)
    return result

print(f"Level-order (BFS): {level_order(root)}")
# Ожидаемый вывод: ['F', 'B', 'G', 'A', 'D', 'I', 'C', 'E', 'H']

Сложность и выбор обхода

Бинарные деревья поиска (BST)

Важное замечание о дубликатах

Классическое определение BST часто не допускает дубликатов ключей. Если дубликаты нужны, их обработку нужно четко определить: либо игнорировать, либо хранить счетчик, либо последовательно помещать их только в одно из поддеревьев (например, в правое, используя >=). Для простоты мы будем придерживаться варианта без дубликатов.

      50      <-- Корень
    /    \
  30      70    <-- (30 < 50), (70 > 50)
 /  \    /  \
20  40  60  80  <-- (20<30), (40>30); (60<70), (80>70)

Поиск в BST

def search_bst(node, target):
    """Ищет target в BST (рекурсивно). Возвращает узел или None."""
    if node is None or node.key == target: # Базовые случаи: не нашли или нашли
        return node
    elif target < node.key:
        return search_bst(node.left, target) # Идем влево
    else: # target > node.key
        return search_bst(node.right, target) # Идем вправо

# --- Создаем BST ---
#       50
#      /  \
#    30    70 ... и т.д. (как в пред. примере)
root_bst = TreeNode(50)
root_bst.left = TreeNode(30); root_bst.right = TreeNode(70)
root_bst.left.left = TreeNode(20); root_bst.left.right = TreeNode(40)
root_bst.right.left = TreeNode(60); root_bst.right.right = TreeNode(80)

key_to_find = 40
found_node = search_bst(root_bst, key_to_find)
print(f"Поиск {key_to_find}: {'Найден' if found_node else 'Не найден'}")

key_to_find = 99
found_node = search_bst(root_bst, key_to_find)
print(f"Поиск {key_to_find}: {'Найден' if found_node else 'Не найден'}")

Проблема баланса

   10
    \
     20
      \
       30
        \
         40
          \
           50

Вставка в BST

def insert_bst(node, key_to_insert):
    """Вставляет key_to_insert в BST (рекурсивно). Возвращает корень поддерева."""
    # Базовый случай: нашли место, создаем узел
    if node is None:
        return TreeNode(key_to_insert)

    # Идем влево или вправо (без вставки дубликатов)
    if key_to_insert < node.key:
        node.left = insert_bst(node.left, key_to_insert)
    elif key_to_insert > node.key:
        node.right = insert_bst(node.right, key_to_insert)
    # else: ключ уже есть, ничего не делаем

    return node # Возвращаем узел (возможно, с обновленным потомком)

# --- Вспомогательная функция для печати (In-order) ---
def get_inorder_keys(node):
    return get_inorder_keys(node.left) + [node.key] + get_inorder_keys(node.right) if node else []

bst_root = None
keys_to_insert = [50, 30, 70, 20, 40, 60, 80, 35]
print(f"--- Вставка в BST ---") # Оставим для ясности
print(f"Вставляем ключи: {keys_to_insert}")
for key in keys_to_insert:
    bst_root = insert_bst(bst_root, key) # Переприсваиваем корень

print(f"Результат (In-order): {get_inorder_keys(bst_root)}")
# [20, 30, 35, 40, 50, 60, 70, 80]

      50
     /  \
   30    70
  / \   / \
 20  40 60  80
     /
    35

Удаление из BST

Кучи и модуль heapq

import heapq

# Очередь задач: (приоритет, описание). Меньше число = Высший приоритет.
tasks_queue = []

heapq.heappush(tasks_queue, (5, "Low priority task"))    # Низкий приоритет
heapq.heappush(tasks_queue, (1, "Urgent bugfix"))   # Высший приоритет!
heapq.heappush(tasks_queue, (3, "Write documentation")) # Средний приоритет
heapq.heappush(tasks_queue, (1, "Reply to boss"))     # Тоже высший приоритет

print(f"Текущая задача (высший приоритет, без извлечения): {tasks_queue[0]}") # O(1)

print("Обработка задач по приоритету:")
while tasks_queue:
    priority, task = heapq.heappop(tasks_queue) # Извлекаем с наименьшим приоритетом O(log n)
    print(f" - Приоритет {priority}: Выполняем '{task}'")

Заключение: деревья покорены!

Деревья — это мощнейший инструмент для представления иерархий и организации данных. Бинарные деревья поиска предлагают эффективный способ работы с упорядоченными данными (O(log n) в среднем), а кучи незаменимы для быстрого доступа к минимальному/максимальному элементу. Понимание различных алгоритмов обхода (DFS, BFS) является ключом к решению множества задач на деревьях.