Реализация Trie (префиксное дерево)

Реализация Trie (префиксное дерево)

Trie — это древовидная структура данных для эффективного хранения и поиска строк, особенно когда нужно искать по префиксам. Каждый узел содержит ссылки на дочерние узлы для каждого символа алфавита.

Структура Trie

Вставляем слова: "apple", "app", "apricot"

        root
        /
       a
       |
       p
      / \
     p   r
     |   |
     l   i
     |   |
     e   c
    (★) |
         o
         |
         t
        (★)

Узлы с (★) отмечены как конец слова

Основная реализация

class TrieNode:
    """Узел префиксного дерева"""
    def __init__(self):
        self.children = {}  # Ссылки на потомков {a: TrieNode, b: TrieNode, ...}
        self.is_end_of_word = False  # Флаг: конец ли слова в этом узле

class Trie:
    """Префиксное дерево для быстрого поиска по префиксам"""
    
    def __init__(self):
        self.root = TrieNode()
    
    def insert(self, word):
        """
        Вставляет слово в Trie
        
        Сложность: O(m), где m — длина слова
        """
        node = self.root
        
        for char in word:
            # Если символа нет, создаём новый узел
            if char not in node.children:
                node.children[char] = TrieNode()
            
            # Переходим к следующему узлу
            node = node.children[char]
        
        # Отмечаем конец слова
        node.is_end_of_word = True
    
    def search(self, word):
        """
        Проверяет, существует ли точное слово в Trie
        
        Сложность: O(m), где m — длина слова
        """
        node = self._find_node(word)
        return node is not None and node.is_end_of_word
    
    def startsWith(self, prefix):
        """
        Проверяет, есть ли слова с данным префиксом
        
        Сложность: O(m), где m — длина префикса
        """
        return self._find_node(prefix) is not None
    
    def _find_node(self, s):
        """
        Вспомогательный метод для поиска узла по строке
        Возвращает узел или None, если строка не найдена
        """
        node = self.root
        
        for char in s:
            if char not in node.children:
                return None
            node = node.children[char]
        
        return node

Пошаговое выполнение

1. Создаём Trie
   trie = Trie()
   Структура: root

2. Вставляем "apple"
   trie.insert("apple")
   
   Шаг за шагом:
   - a: создаём узел, переходим
   - p: создаём узел, переходим
   - p: создаём узел, переходим
   - l: создаём узел, переходим
   - e: создаём узел, переходим
   - Отмечаем is_end_of_word = True
   
   Структура:
   root → a → p → p → l → e (★)

3. Ищем "apple"
   trie.search("apple")
   
   Идём по пути: root → a → p → p → l → e
   Узел существует И is_end_of_word = True
   Результат: True

4. Ищем "app"
   trie.search("app")
   
   Идём по пути: root → a → p → p
   Узел существует, но is_end_of_word = False
   Результат: False

5. Проверяем префикс "app"
   trie.startsWith("app")
   
   Идём по пути: root → a → p → p
   Узел существует
   Результат: True

Полная реализация с расширенной функциональностью

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end_of_word = False
        self.word = None  # Сохраняем само слово для удобства

class Trie:
    def __init__(self):
        self.root = TrieNode()
    
    def insert(self, word):
        """
        Вставляет слово в Trie: O(m)
        """
        if not word:
            return
        
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        
        node.is_end_of_word = True
        node.word = word
    
    def search(self, word):
        """
        Точный поиск слова: O(m)
        """
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]
        
        return node.is_end_of_word
    
    def startsWith(self, prefix):
        """
        Поиск по префиксу: O(m)
        """
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        
        return True
    
    def get_words_with_prefix(self, prefix):
        """
        Возвращает все слова с данным префиксом
        """
        node = self.root
        
        # Находим узел префикса
        for char in prefix:
            if char not in node.children:
                return []
            node = node.children[char]
        
        # DFS для сбора всех слов
        result = []
        self._dfs(node, result)
        return result
    
    def _dfs(self, node, result):
        """
        Рекурсивный обход дерева для сбора слов
        """
        if node.is_end_of_word:
            result.append(node.word)
        
        for child in node.children.values():
            self._dfs(child, result)
    
    def delete(self, word):
        """
        Удаляет слово из Trie
        """
        def _delete_helper(node, word, index):
            if index == len(word):
                if not node.is_end_of_word:
                    return False
                node.is_end_of_word = False
                return len(node.children) == 0
            
            char = word[index]
            if char not in node.children:
                return False
            
            should_delete_child = _delete_helper(
                node.children[char], word, index + 1
            )
            
            if should_delete_child:
                del node.children[char]
                return len(node.children) == 0 and not node.is_end_of_word
            
            return False
        
        _delete_helper(self.root, word, 0)

# Тестирование
if __name__ == "__main__":
    trie = Trie()
    
    # Вставка
    words = ["apple", "app", "apricot", "banana", "band", "can"]
    for word in words:
        trie.insert(word)
    
    # Поиск слов
    print("Поиск слов:")
    print(f"search(apple): {trie.search(apple)}")  # True
    print(f"search(app): {trie.search(app)}")      # True
    print(f"search(ap): {trie.search(ap)}")        # False
    print(f"search(banana): {trie.search(banana)}") # True
    print(f"search(ban): {trie.search(ban)}")      # False
    
    # Поиск по префиксу
    print("\nПоиск по префиксу:")
    print(f"startsWith(app): {trie.startsWith(app)}")      # True
    print(f"startsWith(ban): {trie.startsWith(ban)}")      # True
    print(f"startsWith(xyz): {trie.startsWith(xyz)}")      # False
    
    # Получение слов с префиксом
    print("\nСлова с префиксом:")
    print(f"Слова на ap: {trie.get_words_with_prefix(ap)}")    # [apple, app, apricot]
    print(f"Слова на ban: {trie.get_words_with_prefix(ban)}")  # [banana, band]
    print(f"Слова на ba: {trie.get_words_with_prefix(ba)}")    # [banana, band]

Анализ сложности

Где m — длина слова, n — количество слов с префиксом.

Когда использовать Trie

✅ Автодополнение (асстсм, IDE) ✅ Проверка орфографии ✅ IP маршрутизация (longest prefix matching) ✅ Словари и T9 на мобильных ❌ Поиск по полному совпадению — лучше хешсет

Альтернатива: Ternary Search Tree

Тернарное дерево поиска экономит память, особенно для больших алфавитов, но медленнее на поиск.

Trie — идеальный выбор для работы с префиксами и автодополнением благодаря линейной сложности от длины слова.