Shortest unique prefix for each word: Frontier Interview's solution

Question: Shortest unique prefix for each word

Companies: Meta

Difficulty: Medium

Problem statement

Given a list of distinct words, return for each word its shortest unique prefix: the shortest leading substring of the word that no other word in the list starts with. Return the answers in the same order as the input.

def shortest_unique_prefixes(words: list[str]) -> list[str]: ...

Constraints

1 ≤ len(words) ≤ 10⁵
1 ≤ len(words[i]) ≤ 100
Words contain only lowercase letters a–z.
No word is a prefix of another word, so every word has a unique prefix.

Example 1

Input: ["kiwi", "kale", "kelp", "fig"]
Output: ["ki", "ka", "ke", "f"]

All three k words share k, but their second letters differ. fig is the only word starting with f.

Example 2

Input: ["tram", "trap", "tree", "zoo"]
Output: ["tram", "trap", "tre", "z"]

tram and trap first differ at the fourth letter, so each needs the whole word.

Approach

A prefix of word w identifies w exactly when no other word starts with that
prefix. So we need, for each word, the first position i where the prefix
w[:i+1] is shared by exactly one word: w itself.

Brute force. For each word, try prefixes of length 1, 2, 3, ... and scan
every other word to see if it shares that prefix. That is O(n² · L) in the
worst case. Fine to mention, too slow for 10⁵ words.

Trie with pass-through counts (optimal). Insert every word into a trie.
Each node stores count, the number of words whose path passes through it.
Then walk each word again from the root. The first node with count == 1 ends
the shortest unique prefix: only this word reaches that node, so no other word
starts with that prefix, and every shorter prefix had count >= 2, so it was
shared.

Invariant: after all insertions, node.count equals the number of input words
that have the node's string as a prefix. Counts never increase as you go deeper,
so the first 1 on the path is the answer.

If a walk reaches the end of the word without seeing count == 1, the word is
a prefix of another word or a duplicate. The stated constraints rule that out,
but the code returns the whole word in that case instead of crashing (that is
the natural answer to discuss as a follow-up).

Sort alternative. Sort the words. A word's shortest unique prefix is one
character longer than its longest common prefix with either sorted neighbor,
since the neighbors are the words most similar to it. This uses O(n) extra
memory instead of a trie, at the cost of O(n log n) comparisons. Good to
mention as a memory-friendly variant.

Code

from __future__ import annotations


class TrieNode:
    __slots__ = ("children", "count")

    def __init__(self) -> None:
        self.children: dict[str, TrieNode] = {}
        self.count = 0  # number of words whose path passes through this node


def shortest_unique_prefixes(words: list[str]) -> list[str]:
    """Return the shortest unique prefix of each word, in input order."""
    root = TrieNode()
    for word in words:
        node = root
        for ch in word:
            node = node.children.setdefault(ch, TrieNode())
            node.count += 1

    result: list[str] = []
    for word in words:
        node = root
        length = len(word)  # fallback: word is a prefix of another or a duplicate
        for i, ch in enumerate(word):
            node = node.children[ch]
            if node.count == 1:
                length = i + 1
                break
        result.append(word[:length])
    return result


def shortest_unique_prefixes_sorted(words: list[str]) -> list[str]:
    """Same answer using sorting and neighbor LCPs instead of a trie."""

    def lcp(a: str, b: str) -> int:
        k = 0
        while k < len(a) and k < len(b) and a[k] == b[k]:
            k += 1
        return k

    order = sorted(range(len(words)), key=lambda i: words[i])
    result = [""] * len(words)
    for pos, idx in enumerate(order):
        word = words[idx]
        need = 0
        if pos > 0:
            need = max(need, lcp(word, words[order[pos - 1]]))
        if pos + 1 < len(order):
            need = max(need, lcp(word, words[order[pos + 1]]))
        result[idx] = word[: min(need + 1, len(word))]
    return result


if __name__ == "__main__":
    cases = [
        (["planet", "plane", "pluto", "mars"], None),  # 'plane' is a prefix
        (["kiwi", "kale", "kelp", "fig"], ["ki", "ka", "ke", "f"]),
        (["apple"], ["a"]),
        (["ab", "ac", "b"], ["ab", "ac", "b"]),
        (["tram", "trap", "tree", "zoo"], ["tram", "trap", "tre", "z"]),
    ]
    for words, expected in cases:
        a = shortest_unique_prefixes(words)
        b = shortest_unique_prefixes_sorted(words)
        assert a == b, (words, a, b)
        if expected is not None:
            assert a == expected, (words, a)
    assert shortest_unique_prefixes(["planet", "plane", "pluto", "mars"]) == [
        "planet", "plane", "plu", "m"
    ]
    print("all tests passed")

Complexity

Let n be the number of words and S the total number of characters.

Trie build: O(S) time, O(S) space for the nodes in the worst case
(each node holds a small dict; with a fixed 26-slot array it is O(26 · S)).
Query pass: O(S) time in the worst case, and in practice only the length
of each answer.
Total: O(S) time, O(S) space, plus O(S) for the output.
Sort variant: O(S log n) time for the comparisons, O(n) extra space
beyond the output.

Test cases

["kiwi", "kale", "kelp", "fig"] → ["ki", "ka", "ke", "f"]
["tram", "trap", "tree", "zoo"] → ["tram", "trap", "tre", "z"]
(tram and trap differ only at the last letter, so the prefix is the whole word)
["apple"] → ["a"] (a single word needs one letter)
["ab", "ac", "b"] → ["ab", "ac", "b"]
["a", "b", "c"] → ["a", "b", "c"]
Output order must match input order: ["mint", "melon", "mango"] → ["mi", "me", "ma"]
Outside the constraints: ["plane", "planet"] → ["plane", "planet"]
(planet is unique at its last letter, but no prefix of plane is unique;
we return the whole word, which still matches planet; discuss how to
report that)
Large input: 10⁵ words of length 100 sharing a 99-letter prefix runs in
about 10⁷ steps.

Follow-ups

A word can be a prefix of another, or appear twice. Store an
end_count on terminal nodes. A word that ends on a node with
count > 1 has no unique prefix; return None or the full word plus a
flag, as the interviewer prefers. Duplicates get the same treatment.
Words arrive and leave over time. Keep the trie live. insert
increments counts along the path; delete decrements them and prunes
nodes that reach 0. A query is one walk down the word, O(L). Answers for
other words change only along the path you touched.
Memory is tight (10⁵ long words, many shared paths). Use the sorted
neighbor-LCP method, or a compressed trie (radix tree) that stores edge
labels as (word_index, start, end) slices. For a static list, the sort
method needs no trie at all.
Shortest unique prefix for autocomplete or command abbreviation
(like git accepting git sta). Precompute the answers once; at lookup
time, walk the typed prefix and accept it if the node count is 1. Mention
the case-insensitive or Unicode alphabet change: dict children handle it
without code changes.
Shortest unique substring anywhere in the word, not just a prefix.
Much harder: build a generalized suffix automaton or suffix array over
all words and find, for each word, the shortest substring that occurs in
only that word. Expect to describe it, not code it.

Interviewer expectations

Mid-level (L4)

Sees the trie quickly and explains the per-node count.
Writes working insert and query code that keeps input order.
States O(total characters) time and space correctly.

Senior (L5)

Starts with the brute force, then justifies why the first count == 1
node is correct (counts never increase down a path).
Handles the out-of-constraint cases (prefix words, duplicates) on purpose,
not by accident.
Offers the sort-based alternative and compares memory against the trie.

Staff (L6+)

Discusses the dynamic version (insert/delete with pruning) and what it
means for previously returned answers.
Reasons about memory layout at scale: dict vs array children, compressed
tries, or no trie at all.
Connects the problem to real uses (CLI abbreviations, ID shortening like
short git hashes) and the guarantees those need.