Implement Trie Prefix Tree

The problem

Build a word store called Trie. You get a list operations of operation names and a list values with the arguments of each operation, in the same order. Carry the operations out in order and return a list with one result for each. "Trie" comes first and creates an empty store. "insert" adds its word; adding a word that is already stored changes nothing. "search" gives true if exactly that word was inserted, otherwise false. "startsWith" gives true if some inserted word begins with the given prefix, otherwise false; a whole word counts as a prefix of itself. "Trie" and "insert" give null.

What to write

function solve(operations: string[], values: string[][]): (boolean | null)[]

Input
operations = ["Trie","insert","search","search","startsWith","insert","search"]
values = [[],["apple"],["apple"],["app"],["app"],["app"],["app"]]

Output
[null,null,true,false,true,null,true]

Input
operations = ["Trie","search","startsWith"]
values = [[],["a"],["a"]]

Output
[null,false,false]

Constraints

Words contain lowercase English letters.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   A plain list of the inserted words works, but then every search and every startsWith may have to compare against each word in the list. Aim to make the work depend only on the length of the word you are asked about, not on how many words are stored.

2. Hint 2: Ask a smaller question

   Say apple, apply and banana are stored, and you are asked whether any word starts with "ap". Read "ap" one letter at a time. After "a", which words are still possible? After "ap"? What do apple and apply have in common?

3. Hint 3: Name the pattern

   Questions about how words begin are the signal for a trie, also called a prefix tree. It is a tree where each step down is one letter, so words that begin the same share their first steps, and a word is a path from the top. insert, search and startsWith each walk down one word's path. Read the lesson on trie.

4. Hint 4: State the invariant

   Each stop on a path, called a node, stands for the letters on the path from the top to it. Below the top, a node exists only if some inserted word begins with those letters. A node is marked only if an inserted word ends exactly there. So search needs path and mark; startsWith needs only the path.

5. Hint 5: Sketch the pseudocode

   Root: no children, not marked. insert, search and startsWith start at the root and go down one letter at a time. insert creates each missing child and marks the last node. search and startsWith answer false if a child is missing; otherwise startsWith is true, and search is true if the last node is marked. Trie and insert give null.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Share prefixes through character nodes and mark complete words separately from prefixes.

The slow way

Keep the inserted words in a list. search looks for the word in the list, and startsWith checks each stored word for the prefix. It is correct, but a question may have to compare against every stored word, and words that begin the same have their first letters stored and compared again and again.

The insight

Words that begin the same repeat the same first letters. A trie, or prefix tree, stores each shared beginning once, as a path of letters from a single starting point, and marks where each word ends. A question about a word then only follows that word's letters, so the work depends on that word's length, not on how many words are stored.

Step by step

1. Start with an empty root node: no children, not the end of a word. A node is one stop on a path of letters. It keeps its children, one for each letter that can come next, and a mark for the end of a word. The root is the stop before any letter.
2. To insert, follow the word's letters, adding each node that is missing. Each letter is one step down. A node that already exists is reused, which is how words with the same beginning share their first steps.
3. Mark the last node as the end of a word. A path alone cannot tell a stored word from the beginning of one. After apple is inserted the path for app exists, and only a mark can say whether app itself was inserted.
4. To look up a word or prefix, follow its letters; a missing one fails. A look-up only reads; it never adds nodes. A letter with no child means no inserted word begins with the letters read so far, so the look-up ends with no result.
5. For search, the last node must be marked as the end of a word. The path can exist for a mere beginning, such as app after only apple was inserted. The mark is what says that exactly this word was inserted.
6. For startsWith, finding every letter of the prefix is enough. A node below the root exists only because some inserted word passes through it, so reaching the node of a prefix's last letter proves a word begins with the prefix. The mark does not matter here.
7. Carry out the operations in order and collect one result for each. The operations arrive as one list, so the answer is a list too: null for Trie and insert, true or false for each search and startsWith.

The same numbered steps are the comments in each solution.

function solve(operations, values) {
  // 1. Start with an empty root node: no children, not the end of a word
  const root = { end: false, children: new Map() };

  // 2. To insert, follow the word's letters, adding each node that is missing
  function insert(word) {
    let node = root;
    for (const letter of word) {
      if (!node.children.has(letter)) {
        node.children.set(letter, { end: false, children: new Map() });
      }
      node = node.children.get(letter);
    }
    // 3. Mark the last node as the end of a word
    node.end = true;
  }

  // 4. To look up a word or prefix, follow its letters; a missing one fails
  function find(prefix) {
    let node = root;
    for (const letter of prefix) {
      if (!node.children.has(letter)) return null;
      node = node.children.get(letter);
    }
    return node;
  }

  // 5. For search, the last node must be marked as the end of a word
  function search(word) {
    const node = find(word);
    return node !== null && node.end;
  }

  // 6. For startsWith, finding every letter of the prefix is enough
  function startsWith(prefix) {
    return find(prefix) !== null;
  }

  // 7. Carry out the operations in order and collect one result for each
  const out = [];
  for (let i = 0; i < operations.length; i++) {
    const op = operations[i];
    const args = values[i];
    let result = null;
    if (op === "insert") {
      insert(args[0]);
    } else if (op === "search") {
      result = search(args[0]);
    } else if (op === "startsWith") {
      result = startsWith(args[0]);
    }
    out.push(result);
  }
  return out;
}

function solve(operations: string[], values: string[][]): (boolean | null)[] {
  // A node keeps its children by letter and whether a word ends at it.
  type Node = { end: boolean; children: Map<string, Node> };

  // 1. Start with an empty root node: no children, not the end of a word
  const root: Node = { end: false, children: new Map() };

  // 2. To insert, follow the word's letters, adding each node that is missing
  function insert(word: string): void {
    let node = root;
    for (const letter of word) {
      if (!node.children.has(letter)) {
        node.children.set(letter, { end: false, children: new Map() });
      }
      node = node.children.get(letter)!;
    }
    // 3. Mark the last node as the end of a word
    node.end = true;
  }

  // 4. To look up a word or prefix, follow its letters; a missing one fails
  function find(prefix: string): Node | null {
    let node = root;
    for (const letter of prefix) {
      if (!node.children.has(letter)) return null;
      node = node.children.get(letter)!;
    }
    return node;
  }

  // 5. For search, the last node must be marked as the end of a word
  function search(word: string): boolean {
    const node = find(word);
    return node !== null && node.end;
  }

  // 6. For startsWith, finding every letter of the prefix is enough
  function startsWith(prefix: string): boolean {
    return find(prefix) !== null;
  }

  // 7. Carry out the operations in order and collect one result for each
  const out: (boolean | null)[] = [];
  for (let i = 0; i < operations.length; i++) {
    const op = operations[i];
    const args = values[i];
    let result: boolean | null = null;
    if (op === "insert") {
      insert(args[0]);
    } else if (op === "search") {
      result = search(args[0]);
    } else if (op === "startsWith") {
      result = startsWith(args[0]);
    }
    out.push(result);
  }
  return out;
}

def solve(operations, values):
    # 1. Start with an empty root node: no children, not the end of a word
    root = {"end": False, "children": {}}

    # 2. To insert, follow the word's letters, adding each node that is missing
    def insert(word):
        node = root
        for letter in word:
            if letter not in node["children"]:
                node["children"][letter] = {"end": False, "children": {}}
            node = node["children"][letter]
        # 3. Mark the last node as the end of a word
        node["end"] = True

    # 4. To look up a word or prefix, follow its letters; a missing one fails
    def find(prefix):
        node = root
        for letter in prefix:
            if letter not in node["children"]:
                return None
            node = node["children"][letter]
        return node

    # 5. For search, the last node must be marked as the end of a word
    def search(word):
        node = find(word)
        return node is not None and node["end"]

    # 6. For startsWith, finding every letter of the prefix is enough
    def starts_with(prefix):
        return find(prefix) is not None

    # 7. Carry out the operations in order and collect one result for each
    out = []
    for op, args in zip(operations, values):
        result = None
        if op == "insert":
            insert(args[0])
        elif op == "search":
            result = search(args[0])
        elif op == "startsWith":
            result = starts_with(args[0])
        out.append(result)
    return out

Example 1, traced

Input
operations = ["Trie","insert","search","search","startsWith","insert","search"]
values = [[],["apple"],["apple"],["app"],["app"],["app"],["app"]]

Output
[null,null,true,false,true,null,true]

Why it works

Every node stands for the letters on the path to it, and a node below the root exists only if some inserted word begins with those letters. So reaching the node of a prefix's last letter proves a word begins with it. The mark is set only where an inserted word ends, so a marked last node proves that exact word was inserted.

Time and space

Time O(word length) per operation, extra space O(total inserted characters). Each operation follows one node per letter of its word, and each step is one look-up among at most 26 children: O(word length) time, however many words are stored. The trie has at most one node per inserted letter, plus the root: O(total inserted characters) space.

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• Use the mark for search only. Without it, search app answers true when only apple was inserted. If startsWith needs the mark too, startsWith app answers false in that same state.
• Reuse a node that already exists. If insert makes a fresh node for every letter, inserting app after apple replaces the path that apple uses, and search apple then answers false.
• Do not add nodes while looking up. If a look-up creates the missing letters as it walks, search "a" on an empty store leaves a node behind, and startsWith "a" then answers true, although no word was inserted.

Pattern

Trie. Builds on Two Sum and Maximum Depth of Binary Tree.

Keep going