Implement Trie (Prefix Tree)

A prefix tree (also known as a trie) is a tree data structure used to efficiently store and retrieve keys in a set of strings. Some applications of this data structure include auto-complete and spell checker systems.

Implement the PrefixTree class:

• PrefixTree() Initializes the prefix tree object.
• void insert(String word) Inserts the string word into the prefix tree.
• boolean search(String word) Returns true if the string word is in the prefix tree (i.e., was inserted before), and false otherwise.
• boolean startsWith(String prefix) Returns true if there is a previously inserted string word that has the prefix prefix, and false otherwise.

PrefixTree prefixTree = new PrefixTree();
prefixTree.insert("dog");
prefixTree.search("dog"); // return true
prefixTree.search("do"); // return false
prefixTree.startsWith("do"); // return true
prefixTree.insert("do");
prefixTree.search("do"); // return true

Approach

Build a trie node with a map of children and an end-of-word flag. For insert, walk the path creating nodes as needed and mark the end. For search, walk the path and return true only if the end flag is set. For startsWith, walk the prefix path and return true if it exists.

Why It Works

Each word corresponds to a path from the root through its characters, and the end flag marks complete words. Search succeeds exactly when that path exists and ends at a word boundary. Prefix checks only require that the path exists, which is exactly what startsWith tests.

Complexity

• Time: O(L) per operation, where L is the word/prefix length
• Space: O(L) extra for each inserted word (total trie size)

Reference Implementation (Python)

class TrieNode:
    def __init__(self):
        self.children = {}
        self.end = False

class Trie:

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

    def insert(self, word: str) -> None:
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.end = True

    def search(self, word: str) -> bool:
        node = self._traverse(word)
        return node is not None and node.end

    def startsWith(self, prefix: str) -> bool:
        return self._traverse(prefix) is not None

    def _traverse(self, s: str):
        node = self.root
        for ch in s:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node