Недавно увидел на просторах телеграмма заметку о том как решать алгоритмические задачи на деревья. Вспомнил, что в свое время у меня тоже были некоторые наработки, при этом они отличаются от того что описывается по умолчанию в статьях и курсах. Решил "показать их миру" и похоливарить, поэтому очень рассчитываю на ваш фидбэк, вдруг вам эта информация пригодится во время подготовки к собеседованиям. А может быть узнаете что-то новое для себя.

Введение. Что такое двоичное дерево поиска?

Двоичное дерево поиска (англ. binary search tree, BST) — двоичное дерево, для которого выполняются следующие дополнительные условия (свойства дерева поиска):

• оба поддерева — левое и правое — являются двоичными деревьями поиска;

• у всех узлов левого поддерева произвольного узла X значения ключей данных меньше либо равны, нежели значение ключа данных самого узла X;

• у всех узлов правого поддерева произвольного узла X значения ключей данных больше, нежели значение ключа данных самого узла X.

Источник

Пример дерева для которого соблюдаются все указанные свойства:

Чем полезно дерево поиска?

• В нем можно быстро искать элементы за O(log(n))

• Его можно обходить разными способами.

Рассмотрим 3 варианта обхода деревьев.

Первые попытки обойти дерево

Хватит теории, пора кодить. На Leetcode есть целая задача 94. Binary Tree Inorder Traversal. Данный алгоритм обхода дерева полезен если перед нами стоит задача посетить узлы дерева в порядке возрастания. Суть задачи - вернуть массив значений всех узлов дерева в порядке возрастания.

class Solution:
    def inorderTraversal(self, root):
        acc = []
        
        self.helper(root, acc)
        
        return acc
    
    def helper(self, root, acc):
        if not root:
            return
        
        self.helper(root.left, acc)
        acc.append(root.val)
        self.helper(root.right, acc)

class Solution:
    def inorderTraversal(self, root):
        stack = []
            
        curr = root
        acc = []
        while(stack or curr):
            while(curr):
                stack.append(curr)
                curr = curr.left
            
            if stack:
                curr = stack.pop()
                acc.append(curr.val)
                curr = curr.right
        
        return acc

Обходим дерево "в стиле Python"

Казалось бы, целых 2 решения придумано, стоит ли дальше заморачиваться. Я увидел что дальше LeetCode предлагает мне закодить 2 остальных варианта обхода:

• pre-order (NLR)

• post-order(LRN)

Пришла мысль, можно ли к этим задачам применить более творческий подход чем просто копипастить одно и тоже. Вспомнил букву S из SOLID и возникла идея - отделить логику обхода дерева от остального кода. Ведь на литкоде есть не только задачи на сбор узлов в массив. Вот что получилось:

class Solution:
    def inorderTraversal(self, root):
        return [node.val for node in self.inOrder(root)]

    def inOrder(self, node):
        if not node:
            return
        
        yield from self.inOrder(node.left)
        yield node
        yield from self.inOrder(node.right)

Я воспользовался конструкциями yield и yield from чтобы организовать ленивый перебор дерева в нужном порядке.

Вдохновленный пошел набивать себе рейтинг на задачах 144. Binary Tree Preorder Traversal и 145. Binary Tree Postorder Traversal

class Solution:
    def preorderTraversal(self, root):
        return [node.val for node in self.preOrder(root)]


    def preOrder(self, node):
        if not node:
            return
        
        yield node
        yield from self.preOrder(node.left)
        yield from self.preOrder(node.right)     

class Solution:
    def postorderTraversal(self, root):
        return [node.val for node in self.postOrder(root)]

    def postOrder(self, node):
        if not node:
            return
        
        yield from self.postOrder(node.left)
        yield from self.postOrder(node.right)
        yield node

+3 решенных задачки в копилку, но все равно было ощущение что это все ерунда, с таким же успехом можно было бы просто адаптировать рекурсивное решение для всех 3х случаев. Не был очевиден профит от генераторов.

Усложняем задачу. Обходим дерево в 2 указателя.

Спустя время мне попадается задача 653. Two Sum IV - Input is a BST. Смысл задачи - вернуть 2 числа хранящихся в дереве сумма которых равна заданному числу.

До того как она мне встретилась я решил все 3 предыдущих версии задачи Two Sum и вспомнил алгоритм решения для ситуации когда массив заранее отсортирован.

И вспомнил что у меня же есть мои генераторы, написанные ранее, с их помощью реализовать обход дерева для поиска двух чисел выглядит "делом техники":

class Solution:
    def findTarget(self, root, k):
        left = self.inOrder(root)
        right = self.postOrder(root)
        

        l = next(left)
        r = next(right)

        while(True):
            try:
                if l + r == k and l != r:
                    return True
                elif l + r < k:
                    l = next(left)
                else:
                    r = next(right)
            except StopIteration:
                return False
    

    def inOrder(self, root):
        if not root:
            return
        
        yield from self.inOrder(root.left)
        yield root.val
        yield from self.inOrder(root.right)


    def postOrder(self, root):
        if not root:
            return
        
        yield from self.postOrder(root.right)
        yield root.val
        yield from self.postOrder(root.left)

Результат - еще одна закрытая задачка. К слову это не единственный случай когда пригодились генераторы, но для демонстрации идеи думаю его хватит. На просторах сети в момент решения я не смог найти ничего подобного, при этом на мой субъективный взгляд красивее и идиоматичнее решить задачу на Python невозможно. Решение осталось лежать и ждать своего часа, и вот спустя время публикую его для всех :)

Вот такая заметка получилась, спасибо что дочитали до конца! Делитесь своими впечатлениями и опытом в комментариях! А если вам после прочтения захотелось ещё - подписывайтесь на мой канал в TG, там я пишу чаще и компактнее и не только про алгоритмы.