← Tries

Drill #107 — Implement Trie / Prefix Tree

Medium Trie Build Tries

In plain English: Build a tree-like structure where each path from root to a node represents a prefix of inserted words. Support inserting words and checking if a word or prefix exists.

Each node is a dictionary of children keyed by character. Insert walks the tree creating nodes; search walks and checks the end marker. Prefix search is just search without the end check.

Prompt

Implement a trie with insert, search, and startsWith methods.

Try to write it from scratch before scrolling down.

Solution

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

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

    def insert(self, word):
        # create missing child nodes along the path for each character
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True

    def search(self, word):
        node = self._find(word)
        # must land on a node AND it must be a complete word, not just a prefix
        return node is not None and node.is_end

    def starts_with(self, prefix):
        return self._find(prefix) is not None

    def _find(self, prefix):
        # traverse character-by-character; return None early if path doesn't exist
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node

# Test: t = Trie(); t.insert("apple")
# t.search("apple") == True; t.search("app") == False
# t.starts_with("app") == True
O(m) time per operation · O(n*m) space total

Related Micro Drills

← Drill #106