Bracket the first matching substring in each word

Question: Bracket the first matching substring in each word

Companies: Meta, Uber

Difficulty: Easy

Problem statement

Implement wrap_substrings(s: str, elements: list[str]) -> str.

s is a sentence of words separated by single spaces. elements is a list of strings. For each word in s, find the first substring of that word that equals one of the elements and wrap it in square brackets. Leave words with no match unchanged, and keep the spaces between words as they are.

Rules for choosing the match in a word:

• Pick the match that starts at the smallest index in the word.
• If several elements match at that index, pick the longest one.
• Wrap only that one match; later matches in the same word are left alone.
• Matching is case-sensitive. Ignore empty strings in elements.

Constraints: 0 <= len(s) <= 10^5, 0 <= len(elements) <= 10^4, each element has at most 100 characters.

Example 1

s = "orange banana kiwi"
elements = ["an", "ana", "wi"]
output: "or[an]ge b[ana]na ki[wi]"

In banana, both an and ana start at index 1, so the longer ana is wrapped.

Example 2

s = "abc hello"
elements = ["c", "ab", "abc"]
output: "[abc] hello"

c also matches in abc, but ab and abc start earlier, and abc is the longer of those. hello has no match.

Start with a straightforward solution, then make it fast when there are many elements.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

Each word is handled on its own. For a word we want the leftmost position where some element starts, and among the elements that start there, the longest one. We wrap that one span in [ and ] and leave the rest of the word alone. Words with no match pass through unchanged, and the single spaces between words are kept.

Brute force. For each word, for each start index i, try every element and check word.startswith(e, i). Stop at the first i with any hit, keeping the longest hit. With W total characters in the sentence, E elements and max element length K, this is O(W * E * K). That is fine for small inputs and is a good first answer to say out loud.

Optimal: a trie of the elements. Put every element in a trie. For each start index i in a word, walk the trie along word[i], word[i+1], ... until the walk falls off. Every node marked as the end of an element is a match starting at i; the deepest one is the longest. The first i that produces any match is the answer for that word, so we stop scanning there.

Why it is correct:

• We visit start indices in increasing order and stop at the first one with a match, so the match is the leftmost.
• At that start, the trie walk sees every element that is a prefix of word[i:], so keeping the deepest end marker gives the longest.
• Each start index costs at most K trie steps, so the scan is O(W * K) and no longer depends on E.

Invariants: the trie only holds non-empty elements (an empty element would "match" everywhere and must be dropped up front); only the first match in a word is wrapped; output words are joined with the same single spaces as the input.

Code

from typing import Dict, List


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

    def __init__(self) -> None:
        self.children: Dict[str, "TrieNode"] = {}
        self.is_end = False


class SubstringWrapper:
    def __init__(self, elements: List[str]) -> None:
        self.root = TrieNode()
        for e in elements:
            if e:  # an empty element would match everywhere
                self._insert(e)

    def _insert(self, e: str) -> None:
        node = self.root
        for ch in e:
            node = node.children.setdefault(ch, TrieNode())
        node.is_end = True

    def _longest_at(self, word: str, i: int) -> int:
        """Length of the longest element that starts at word[i], or 0."""
        node, best = self.root, 0
        for j in range(i, len(word)):
            node = node.children.get(word[j])
            if node is None:
                break
            if node.is_end:
                best = j - i + 1
        return best

    def wrap_word(self, word: str) -> str:
        for i in range(len(word)):
            length = self._longest_at(word, i)
            if length:
                return f"{word[:i]}[{word[i:i + length]}]{word[i + length:]}"
        return word

    def wrap(self, s: str) -> str:
        return " ".join(self.wrap_word(w) for w in s.split(" "))


def wrap_substrings(s: str, elements: List[str]) -> str:
    return SubstringWrapper(elements).wrap(s)


if __name__ == "__main__":
    assert wrap_substrings("orange banana kiwi", ["an", "ana", "wi"]) == "or[an]ge b[ana]na ki[wi]"
    assert wrap_substrings("hello world", ["xyz"]) == "hello world"
    assert wrap_substrings("abc", ["c", "ab", "abc"]) == "[abc]"
    assert wrap_substrings("aaaa", ["aa"]) == "[aa]aa"
    assert wrap_substrings("", ["a"]) == ""
    assert wrap_substrings("cat dog", []) == "cat dog"
    assert wrap_substrings("cat", ["", "t"]) == "ca[t]"
    assert wrap_substrings("Data data", ["da"]) == "Data [da]ta"
    assert wrap_substrings("a  b", ["b"]) == "a  [b]"
    print("ok")

Complexity

Let W be the number of characters in s, E the number of elements, T their total length, and K the longest element.

• Building the trie: O(T) time and O(T) space.
• Wrapping: each start index walks at most K trie nodes, so O(W * K) time in the worst case; usually far less, because the scan stops at the first matching start in each word.
• Output: O(W) extra space for the new string.
• Brute force for comparison: O(W * E * K) time, O(1) extra space.

Test cases

• s = "orange banana kiwi", elements = ["an", "ana", "wi"] → "or[an]ge b[ana]na ki[wi]" (longest wins at the same start).
• s = "hello world", elements = ["xyz"] → "hello world" (no match, unchanged).
• s = "abc", elements = ["c", "ab", "abc"] → "[abc]" (leftmost start beats a later shorter one; longest at that start).
• s = "aaaa", elements = ["aa"] → "[aa]aa" (only the first match per word is wrapped).
• s = "", elements = ["a"] → "".
• s = "cat dog", elements = [] → "cat dog".
• s = "cat", elements = ["", "t"] → "ca[t]" (empty element ignored).
• s = "Data data", elements = ["da"] → "Data [da]ta" (matching is case-sensitive).
• s = "a b", elements = ["b"] → "a [b]" (repeated spaces are kept because we split on a single space).

Follow-ups

1. Wrap every non-overlapping match, not just the first. After wrapping a match at [i, i+len), continue scanning from i + len instead of returning. Still O(W * K).
2. Merge overlapping or adjacent matches into one bracket (LeetCode "Add Bold Tag" style). Mark a boolean per character covered by any match (longest match at each start is enough), then emit [ and ] at the edges of each covered run.
3. Many elements, very long text (streaming or a large corpus). Replace the per-start trie walk with Aho-Corasick: one pass over the text in O(W + matches) after an O(T * alphabet) build. Leftmost-longest needs the standard trick of tracking the earliest start among outputs and confirming the longest one at that start.
4. Priority changes: the longest match anywhere in the word, ties to the earliest. Scan all starts instead of stopping at the first, keep (length, -start) as the best key. Same trie, same bound.
5. Case-insensitive matching or Unicode. Insert casefold() elements, walk with casefolded characters, but slice the original word so the output keeps its case. Mention that casefolding can change length for some characters, so map indices carefully or normalize first.

Interviewer expectations

Mid-level (L4)

• States the tie rules (leftmost, then longest) before coding, or asks about them.
• Writes a correct brute force with startswith, handles no-match words and empty input.
• Gives the right complexity for what they wrote.

Senior (L5)

• Moves to a trie without prompting, explains why it removes the factor of E.
• Clean separation: build once, reuse for many sentences; handles empty elements, repeated spaces, case.
• Adapts quickly to "wrap all matches" and "merge overlaps".

Staff (L6+)

• Discusses Aho-Corasick for large dictionaries and long or streaming text, and how leftmost-longest semantics interact with it.
• Frames the API for production use: precompiled matcher, thread-safe reads, memory cost of the trie versus a hash set of elements bucketed by length.
• Raises product questions: tokenization beyond spaces, punctuation, Unicode normalization, and what to do when matches overlap word boundaries.