Published on

leetcode-208 Implement Trie (Prefix Tree)

Authors
  • avatar
    Name
    Gene Zhang
    Twitter

[208] Implement Trie (Prefix Tree)

Key Concept: Trie Data Structure - A tree where each node represents a character. Used for efficient prefix operations. Each node has up to 26 children (for lowercase letters) and a flag indicating if it's the end of a word.

Use cases: Auto-complete, spell checker, IP routing, dictionary implementation.

# A trie or prefix tree is a tree data structure used to efficiently store and
# retrieve keys in a dataset of strings. Applications include autocomplete and
# spellchecker.
# Implement the Trie class with:
# - Trie() Initializes the trie object.
# - void insert(String word) Inserts the string word into the trie.
# - boolean search(String word) Returns true if word is in the trie.
# - boolean startsWith(String prefix) Returns true if there is a previously
#   inserted string that has the prefix.
#
# Example:
# Input: ["Trie", "insert", "search", "search", "startsWith", "insert", "search"]
#        [[], ["apple"], ["apple"], ["app"], ["app"], ["app"], ["app"]]
# Output: [null, null, true, false, true, null, true]
#
# Constraints:
# 1 <= word.length, prefix.length <= 2000
# word and prefix consist only of lowercase English letters.

class TrieNode:
    def __init__(self):
        self.children = {}  # Dictionary mapping char -> TrieNode
        self.is_end_of_word = False

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        """Insert a word into the trie."""
        node = self.root

        for char in word:
            # Create new node if path doesn't exist
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]

        # Mark end of word
        node.is_end_of_word = True

    def search(self, word: str) -> bool:
        """Returns True if word is in the trie."""
        node = self.root

        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]

        # Must be marked as end of word
        return node.is_end_of_word

    def startsWith(self, prefix: str) -> bool:
        """Returns True if there is any word with the given prefix."""
        node = self.root

        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]

        return True

# Alternative: Using array instead of dict for children
class TrieNode2:
    def __init__(self):
        self.children = [None] * 26  # For 'a' to 'z'
        self.is_end_of_word = False

class Trie2:
    def __init__(self):
        self.root = TrieNode2()

    def insert(self, word: str) -> None:
        node = self.root
        for char in word:
            idx = ord(char) - ord('a')
            if not node.children[idx]:
                node.children[idx] = TrieNode2()
            node = node.children[idx]
        node.is_end_of_word = True

    def search(self, word: str) -> bool:
        node = self.root
        for char in word:
            idx = ord(char) - ord('a')
            if not node.children[idx]:
                return False
            node = node.children[idx]
        return node.is_end_of_word

    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for char in prefix:
            idx = ord(char) - ord('a')
            if not node.children[idx]:
                return False
            node = node.children[idx]
        return True

# Time Complexity:
#   insert: O(m) where m is length of word
#   search: O(m)
#   startsWith: O(m)
# Space Complexity: O(ALPHABET_SIZE * N * M) in worst case
#   where N is number of words, M is average length
#
# Key Pattern: Trie for prefix-based operations
# Much more efficient than storing all strings and checking each one