MediumLeetCode 648Trie — shortest stored prefix

Replace Words (LeetCode 648)

Problem 4 of 10 in Tries — Practice Problems

The problem

You get a list of roots (the dictionary) and a sentence of space-separated words. Any word that starts with a root should be replaced by that root. If several roots match, use the shortest one. Words with no matching root stay as they are. Return the new sentence.

Input:  dictionary = ["cat", "bat", "rat"]
        sentence   = "the cattle was rattled by the battery"
Output: "the cat was rat by the bat"

cattle starts with cat, rattled with rat, and battery with bat. The other words match no root.

How to think about it

Start with brute force. For each word, try every root and check word.startswith(root), keeping the shortest match. With W words and R roots of length up to L, that is O(W·R·L). If the dictionary is large, every word pays for the whole dictionary.

Slightly better: put the roots in a set and, for each word, test its prefixes from shortest to longest. That is O(L) prefixes per word, but slicing and hashing each prefix costs O(L) too, so O(L²) per word.

The key insight: walk the word down a trie of roots, and stop at the first end marker. Insert every root into a trie. Then read a word one character at a time from the root. The first node you reach that marks the end of a root is, by construction, the shortest root that prefixes the word. If you fall off the trie before that, no root matches. Each word now costs O(L) with no repeated work.

Which pattern is this? Trie, prefix lookup. The tell is “replace a word with a prefix from a dictionary”. A handy trick: store the whole root string at its end node (n['$'] = w) so you can return it directly without rebuilding it from the path.

Approach

  1. Build a trie from the dictionary. At each root’s last node, store the root itself under '$'.
  2. For each word in the sentence, start at the trie root and read characters.
  3. If the character is not a child, stop: no root matches, keep the word.
  4. Otherwise step down. If the new node has '$', return the stored root immediately.
  5. Join the results with spaces.

Solution

def replaceWords(dictionary, sentence):
    root = {}
    for w in dictionary:
        n = root
        for c in w: n = n.setdefault(c, {})
        n['$'] = w                          # store the whole root at its end
    def shortest(word):
        n = root
        for c in word:
            if c not in n: break            # fell off the trie: no root
            n = n[c]
            if '$' in n: return n['$']      # first end marker = shortest root
        return word
    return " ".join(shortest(w) for w in sentence.split())
# replaceWords(["cat","bat","rat"], "the cattle was rattled by the battery") -> "the cat was rat by the bat"

Dry run

The trie has three branches from the root: c-a-t, b-a-t, r-a-t. Each word of the sentence:

Word Walk Stopped because Output
the t not at root fell off the
cattle c→a→t '$' at cat cat
was w not at root fell off was
rattled r→a→t '$' at rat rat
by b, then y missing fell off by
battery b→a→t '$' at bat bat

Joined: "the cat was rat by the bat".

Complexity

  • Time: O(D + S), where D is the total length of the dictionary and S the length of the sentence. Each root is inserted once, and each word is read at most once, character by character.
  • Space: O(D) for the trie, plus O(S) for the split words and the output.

Edge cases and pitfalls

  • Shortest root wins. With roots ["catt", "cat"], cattle becomes cat. Returning at the first '$' gives this for free; walking to the deepest match would be wrong.
  • Word shorter than every root (by with root bat): the loop ends without a match and the word is kept.
  • Empty dictionary leaves the sentence unchanged.
  • Set-based version, wrong order. If you test prefixes from a set, go from shortest to longest. Testing longest first returns the longest matching root, which is wrong.