Implement Trie (Prefix Tree)

A trie (prefix tree) stores a set of strings so that both whole words and prefixes can be looked up quickly. Implement a class Trie with these operations:

  • Trie() creates an empty trie.
  • insert(word) adds word to the trie. Inserting a word that is already present changes nothing.
  • search(word) returns true if word was inserted earlier as a complete word, and false otherwise.
  • startsWith(prefix) returns true if at least one inserted word begins with prefix (a word counts as a prefix of itself), and false otherwise.

Each test is a sequence of calls; the expected output lists every call's return value, with null for the constructor and insert.

Example 1

Input: ["Trie","insert","search","search","startsWith","insert","search"] [[],["code"],["code"],["cod"],["cod"],["cod"],["cod"]] Output: [null,null,true,false,true,null,true] Why: After inserting "code", the full word "cod" is not stored, but it is a prefix of "code". Once "cod" is inserted, searching for it succeeds.

Example 2

Input: ["Trie","search","startsWith","insert","insert","startsWith","startsWith","search","startsWith"] [[],["a"],["a"],["team"],["ten"],["te"],["tea"],["tea"],["tex"]] Output: [null,false,false,null,null,true,true,false,false] Why: An empty trie contains nothing. Later, "tea" is a prefix of "team", so startsWith succeeds, but search fails because "tea" itself was never inserted.

Example 3

Input: ["Trie","insert","insert","search","search","startsWith","startsWith"] [[],["map"],["map"],["map"],["maps"],["maps"],["m"]] Output: [null,null,null,true,false,false,true] Why: Inserting the same word twice is harmless; a longer query than any stored word fails.

Constraints

  • 1 <= word.length, prefix.length <= 2000
  • word and prefix contain only lowercase English letters.
  • At most 3 · 104 calls are made in total.

Pattern Path · Design · problem 162 of 169

Storing all words in a hash set makes search fast, but how slow is startsWith then?
Words that share a beginning could share storage for it. Picture a tree where each edge is labelled with one letter and each path from the root spells a prefix.
Give each node a map from letter to child plus a flag meaning "a word ends here". Both queries walk down one letter at a time; they differ only in whether that flag must be set at the end.
Target: O(L) per operation, L = length of the word or prefix time, O(total characters inserted) extra space.

Start simple

Keep a list of words and compare each query against every word: O(N · L) per query for N words of length L.

The key idea

A trie stores each word as a path from the root, one letter per edge. Words with a common beginning share the nodes for that beginning, so every prefix of every inserted word corresponds to exactly one node. Each node keeps its children (a map from letter to node) and a boolean is_word marking that some inserted word ends exactly here.

insert walks from the root, creating any missing child along the way, and sets is_word on the final node. Both queries walk the same way but never create nodes: if a letter has no child, nothing inserted starts with that string, so the answer is false.

If the walk succeeds, the node reached represents the query string. For startsWith that is enough, since the node exists only because some word passes through it. For search the node must also have is_word set; otherwise the query is merely a prefix of a longer word, like "cod" inside "code". Every operation costs time proportional to the length of its string, independent of how many words are stored.

Complexity

  • Time: O(L) per operation, L = length of the word or prefix
  • Space: O(total characters inserted)

Python

class TrieNode:
    def __init__(self):
        self.children = {}    # letter -> TrieNode
        self.is_word = False  # does an inserted word end here?


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.is_word = True

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

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

    def _walk(self, s: str):
        """Follow s from the root; return the node it ends at, or None."""
        node = self.root
        for ch in s:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

JavaScript

class TrieNode {
  constructor() {
    this.children = new Map(); // letter -> TrieNode
    this.isWord = false;       // does an inserted word end here?
  }
}

class Trie {
  constructor() {
    this.root = new TrieNode();
  }

  insert(word) {
    let node = this.root;
    for (const ch of word) {
      if (!node.children.has(ch)) node.children.set(ch, new TrieNode());
      node = node.children.get(ch);
    }
    node.isWord = true;
  }

  search(word) {
    const node = this.walk(word);
    return node !== null && node.isWord;
  }

  startsWith(prefix) {
    return this.walk(prefix) !== null;
  }

  // Follow s from the root; return the node it ends at, or null.
  walk(s) {
    let node = this.root;
    for (const ch of s) {
      node = node.children.get(ch);
      if (node === undefined) return null;
    }
    return node;
  }
}

What the interviewer is checking

It tests whether you can design a small data structure from scratch, choose a clean node representation and explain why prefix queries become independent of the number of stored words.

Mistakes to avoid

  • Forgetting the end-of-word flag, so search returns true for mere prefixes.
  • Creating nodes during search or startsWith, which silently grows the trie.
  • Sharing one mutable dictionary between nodes (e.g. a mutable default argument in Python).

Follow-ups to expect

  • Support deleting a word? (Unset the flag, prune empty nodes.)
  • Return all words with a given prefix? (DFS from the prefix node.)
  • Support '.' as a wildcard letter in search? (Branch over all children.)

In real software

Autocomplete in search boxes and IDEs, spell checkers, and IP routing tables (longest-prefix match) are built on tries or compressed variants of them.

Practice next

Next Post Previous Post
No Comment
Add Comment
comment url