DSA Replacing Words by Roots

The problem

Replacing words by roots swaps each word in a sentence for the shortest dictionary root that opens it. A root is a short stored word that can stand at the start of a longer word.

A word changes only when a root matches its opening letters exactly. When two roots both open a word, the shorter one wins. A word with no matching root stays whole.

cases = [
    (["cat", "dog", "sun"], ["cats", "and", "sunny", "dogs", "chased", "cattle"], "cat and sun dog chased cat"),
    (["cats", "cat"], ["catsup"], "cat"),
    (["dog"], ["zebra", "ran"], "zebra ran"),
]
for roots, words, expected in cases:
    print("roots", roots)
    print("words", words)
    print("expected:", expected)

The second case lists the longer root first. The answer is still cat, so the winner is decided by root length, not by dictionary order. The third case changes nothing at all.

Intuition: take the first root that fits

The first idea tests each root against the word and keeps the first one that fits. It works when the dictionary happens to list short roots before long ones.

Dictionary order is not part of the problem. The same idea fails as soon as a longer root is listed first.

def first_match(roots, word):
    for root in roots:
        if word.startswith(root):
            return root
    return word

print("cattle with [cat, dog, sun]:", first_match(["cat", "dog", "sun"], "cattle"))
print("catsup with [cats, cat]:", first_match(["cats", "cat"], "catsup"))
print("catsup with [cat, cats]:", first_match(["cat", "cats"], "catsup"))

One word, catsup, gets two different answers from two orderings of the same dictionary. The correct answer is cat in both. Any method whose answer depends on list order is answering the wrong question.

Brute force: test every root, keep the shortest fit

The brute force fixes the order bug by testing every root. It keeps the shortest root that opens the word. A counter records the root tests, and every word pays for all of them.

def replace_brute(roots, word, counter):
    best = None
    for root in roots:
        counter[0] += 1
        if word.startswith(root):
            if best is None or len(root) < len(best):
                best = root
    return best if best is not None else word

roots = ["cat", "dog", "sun"]
words = ["cats", "and", "sunny", "dogs", "chased", "cattle"]
counter = [0]
rewritten = [replace_brute(roots, word, counter) for word in words]
print("Sentence:", " ".join(rewritten))
print("root tests:", counter[0])

Six words against three roots cost 18 root tests. The count is the product of the sentence length and the dictionary size. A long sentence with a large dictionary pays that product in full.

The optimal solution: walk the word through a trie of roots

The state is a trie built from the roots, where each node marks whether a root ends there. Walking a word letter by letter and stopping at the first end flag gives the shortest root directly, because shorter roots end earlier on the same path.

A missing letter ends the walk with no root found. The word then stays whole. Reaching the end of the word with no flag seen also leaves the word whole.

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


def build_trie(roots):
    root = TrieNode()
    for word in roots:
        node = root
        for letter in word:
            if letter not in node.children:
                node.children[letter] = TrieNode()
            node = node.children[letter]
        node.is_end = True
    return root


def replace_one(root, word):
    node = root
    for index, letter in enumerate(word):
        if letter not in node.children:
            return word
        node = node.children[letter]
        if node.is_end:
            return word[:index + 1]
    return word

def replace_words(roots, words):
    trie = build_trie(roots)
    return " ".join(replace_one(trie, word) for word in words)

cases = [
    (["cat", "dog", "sun"], ["cats", "and", "sunny", "dogs", "chased", "cattle"], "cat and sun dog chased cat"),
    (["cats", "cat"], ["catsup"], "cat"),
    (["dog"], ["zebra", "ran"], "zebra ran"),
]
for roots, words, expected in cases:
    result = replace_words(roots, words)
    print("roots", roots, "->", result, "| expected:", expected, "| match:", result == expected)

All three cases from the start of the article pass. The dictionary order in the second case no longer matters. The walk meets the end flags in letter order, and letter order is length order along one path.

Walkthrough: cattle through the root trie

The walk uses roots cat, cats and dog. It reads one letter of cattle at a time. After each letter it prints the end flag of the node it reached.

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

def build_trie(roots):
    root = TrieNode()
    for word in roots:
        node = root
        for letter in word:
            if letter not in node.children:
                node.children[letter] = TrieNode()
            node = node.children[letter]
        node.is_end = True
    return root

trie = build_trie(["cat", "cats", "dog"])
node = trie
word = "cattle"
for index, letter in enumerate(word):
    node = node.children[letter]
    print("letter", letter, "- prefix", word[:index + 1], "- root ends here:", node.is_end)
    if node.is_end:
        print("replacement:", word[:index + 1])
        print("letters checked:", index + 1, "of", len(word))
        break

The flag at t ends the walk after three letters. The longer root cats sits further down the same path and is never reached. The last three letters of cattle are never read at all.

Complexity

The brute force tests every root against every word, so its time grows with the product of the dictionary size and the sentence length. The trie walk reads letters only until the first end flag, so each word costs at most its own length in steps.

The trie stores one node per distinct root prefix. Roots that share an opening share nodes.

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

def build_trie(roots):
    root = TrieNode()
    for word in roots:
        node = root
        for letter in word:
            if letter not in node.children:
                node.children[letter] = TrieNode()
            node = node.children[letter]
        node.is_end = True
    return root

def count_nodes(node):
    total = 1
    for child in node.children.values():
        total += count_nodes(child)
    return total

trie = build_trie(["cat", "dog", "sun"])
print("nodes for roots cat, dog, sun:", count_nodes(trie))
for r in [3, 30, 300]:
    print("roots", r, "- brute root tests per word up to", r, "- trie letters checked for cattle:", 3)

Three roots with no shared openings need ten nodes, root included. Against that fixed small cost, the brute test count grows with every root added to the dictionary. The trie cost for cattle stays at three letters.

Edge cases and pitfalls

A word equal to a root is replaced by itself. The walk reaches the end flag on the last letter, and the replacement is the same string.

A word shorter than every root on its path stays whole. The walk runs out of letters before any end flag appears. Confusing that case with a match is the classic pitfall, and the end flag check is what prevents it.

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


def build_trie(roots):
    root = TrieNode()
    for word in roots:
        node = root
        for letter in word:
            if letter not in node.children:
                node.children[letter] = TrieNode()
            node = node.children[letter]
        node.is_end = True
    return root


def replace_one(root, word):
    node = root
    for index, letter in enumerate(word):
        if letter not in node.children:
            return word
        node = node.children[letter]
        if node.is_end:
            return word[:index + 1]
    return word

trie = build_trie(["cat", "dog", "sun"])
print("word equal to a root:", replace_one(trie, "cat"))
print("word shorter than its root:", replace_one(trie, "ca"))
print("empty word:", repr(replace_one(trie, "")))
print("empty sentence:", " ".join(replace_one(trie, w) for w in []))
empty_trie = build_trie([])
print("no roots at all:", replace_one(empty_trie, "cattle"))

The empty word and the empty dictionary both pass through unchanged. No special case handles them. The walk simply never meets an end flag.