0

Merge known bigrams in a sentence, leftmost first

by Frontier Interview· 1 day ago· ⑂ 0 remixes

Question: Merge known bigrams in a sentence, leftmost first

Companies: xAI

Difficulty: Easy

Problem statement

Write a function that splits a sentence into tokens using a dictionary of known two-word phrases (bigrams).

def bigram_split(text: str, bigrams: list[str]) -> list[str]:
    ...
  • text is a non-empty string of lowercase words separated by spaces.
  • bigrams is a list of phrases. Each phrase is two lowercase words separated by one space, for example "machine learning". A bigram is directed: "a b" does not match b a.

Rules:

  1. Read the words from left to right, and keep their original order in the output.
  2. When two adjacent words form a phrase in bigrams, output them as one token, with the two words joined by one space.
  3. Every other word is its own token.
  4. If two possible matches overlap, take the one that starts further left. A word can be part of only one token.

Example 1

text    = "we use machine learning models"
bigrams = ["machine learning", "learning models"]
output  = ["we", "use", "machine learning", "models"]

"machine learning" and "learning models" share the word learning. The leftmost match wins, so models stays on its own.

Example 2

text    = "go go go"
bigrams = ["go go"]
output  = ["go go", "go"]

Constraints

  • 1 <= number of words <= 10^5
  • 0 <= len(bigrams) <= 10^5
  • Words contain only the letters a-z.

Aim for time linear in the length of the input.

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

Approach

The input is a sentence of lowercase words separated by spaces, plus a set of known bigrams such as "new york". Return the token list. Each pair of adjacent words that forms a known bigram becomes one token. Every other word stays on its own. When matches overlap, the leftmost match wins. Words keep their original order.

Clarify first. Ask these before coding:

  • Can the input have extra spaces? Use text.split(), which ignores repeated spaces.
  • Is a dictionary entry always exactly two words? If entries can have extra spaces, normalize them the same way as the text.
  • What does "leftmost wins" mean on a chain? Take a b c with both a b and b c in the dictionary. The answer is ["a b", "c"]. Once b is used, it cannot start another match.

Brute force. For each position, compare the pair at that position with every dictionary entry. That costs O(n * d) string comparisons. The fix is to store the dictionary as a hash set of (first, second) tuples. Each check then costs O(1) expected, plus hashing the two words.

Greedy scan. Keep an index i:

  • If i + 1 < n and (w[i], w[i+1]) is in the set, emit "w[i] w[i+1]" and set i += 2.
  • Otherwise, emit w[i] and set i += 1.

Why this is the leftmost rule. The scan reaches positions in increasing order. At position i, every word before i is already in a token. No earlier unconsumed word is left that could start an overlapping match. So the first match found is the leftmost one still available, and taking it is exactly what the rule asks for. Jumping by 2 drops any match that would overlap it.

Invariants.

  • Every word lands in exactly one token, and the tokens come out in input order. " ".join(tokens).split() gives back the original words.
  • A token is either one word, or two words whose pair is in the dictionary.

Generalize to n-grams. If the dictionary can hold phrases of any length up to L, try the longest length first at each position. This is the leftmost-longest rule, the same rule regex engines and tokenizers use. Index the phrases by their first word, or use a trie of words, so that each position only checks the candidates that can match there.

Code

from typing import Dict, Iterable, List, Set, Tuple


def bigram_split(text: str, bigrams: Iterable[str]) -> List[str]:
    """Merge adjacent words that form a known bigram; leftmost match wins."""
    pairs: Set[Tuple[str, ...]] = {tuple(b.split()) for b in bigrams}
    words = text.split()
    tokens: List[str] = []
    i, n = 0, len(words)
    while i < n:
        if i + 1 < n and (words[i], words[i + 1]) in pairs:
            tokens.append(f"{words[i]} {words[i + 1]}")
            i += 2
        else:
            tokens.append(words[i])
            i += 1
    return tokens


def ngram_split(text: str, phrases: Iterable[str]) -> List[str]:
    """Leftmost-longest match over phrases of any length (follow-up)."""
    by_first: Dict[str, Set[Tuple[str, ...]]] = {}
    lengths: Dict[str, List[int]] = {}
    for p in phrases:
        key = tuple(p.split())
        if len(key) >= 2:
            by_first.setdefault(key[0], set()).add(key)
    for first, keys in by_first.items():
        lengths[first] = sorted({len(k) for k in keys}, reverse=True)

    words = text.split()
    tokens: List[str] = []
    i, n = 0, len(words)
    while i < n:
        step = 1
        for size in lengths.get(words[i], []):  # longest first
            if i + size <= n and tuple(words[i:i + size]) in by_first[words[i]]:
                step = size
                break
        tokens.append(" ".join(words[i:i + step]))
        i += step
    return tokens


def best_weight_split(text: str, weights: Dict[str, float]) -> List[str]:
    """Follow-up: pick phrases (any length) to maximize total weight."""
    table = {tuple(p.split()): w for p, w in weights.items() if len(p.split()) >= 2}
    max_len = max((len(k) for k in table), default=1)
    words = text.split()
    n = len(words)
    best = [0.0] * (n + 1)        # best[i] = max weight for words[i:]
    step = [1] * (n + 1)
    for i in range(n - 1, -1, -1):
        best[i], step[i] = best[i + 1], 1
        for size in range(2, min(max_len, n - i) + 1):
            w = table.get(tuple(words[i:i + size]))
            if w is not None and w + best[i + size] > best[i]:
                best[i], step[i] = w + best[i + size], size
    tokens: List[str] = []
    i = 0
    while i < n:
        tokens.append(" ".join(words[i:i + step[i]]))
        i += step[i]
    return tokens


if __name__ == "__main__":
    d = ["ice cream", "cream cheese", "hot dog"]
    assert bigram_split("i like ice cream cheese", d) == ["i", "like", "ice cream", "cheese"]
    assert bigram_split("hot dog and hot tea", d) == ["hot dog", "and", "hot", "tea"]
    assert bigram_split("solo", d) == ["solo"]
    assert bigram_split("a  b", ["a b"]) == ["a b"]
    assert bigram_split("x x x", ["x x"]) == ["x x", "x"]
    assert bigram_split("cream ice", d) == ["cream", "ice"]
    assert ngram_split("san francisco bay area", ["san francisco", "san francisco bay"]) == [
        "san francisco bay", "area"]
    assert ngram_split("ice cream cheese", d) == ["ice cream", "cheese"]
    assert bigram_split("a b c d", ["a b", "b c", "c d"]) == ["a b", "c d"]
    w = {"a b": 1.0, "b c": 5.0}
    assert bigram_split("a b c", list(w)) == ["a b", "c"]
    assert best_weight_split("a b c", w) == ["a", "b c"]
    assert best_weight_split("a b c d", {"a b c": 2.0, "a b": 1.5, "c d": 1.5}) == ["a b", "c d"]
    assert best_weight_split("solo", w) == ["solo"]
    print("ok")

Complexity

Let n be the number of words, W the total characters in the text, d the number of dictionary entries, and D their total characters.

  • Building the set: O(D) time and space.
  • bigram_split: O(W) time. Each position does at most one set lookup, and hashing a pair costs its length. The output uses O(W) space.
  • ngram_split: at most k lookups per position, where k is the number of distinct phrase lengths that start with that word. Each lookup builds a slice of up to L words, so the worst case is O(n * k * L) word operations. A word trie brings this down to O(n * L), with no slicing.
  • best_weight_split: O(n * L) lookups, each slicing up to L words, so O(n * L^2) word operations. It uses O(n) space for the DP arrays.

Test cases

  • "i like ice cream cheese" with {ice cream, cream cheese} gives ["i", "like", "ice cream", "cheese"]. The leftmost match wins the overlap.
  • "hot dog and hot tea" with {hot dog} gives ["hot dog", "and", "hot", "tea"]. The same first word can still fail to match later.
  • A single word, "solo", gives ["solo"].
  • "x x x" with {x x} gives ["x x", "x"]. A repeated word must not be used twice.
  • "cream ice" with {ice cream} gives ["cream", "ice"]. Order matters, because bigrams are directed.
  • "a b" (two spaces) with {a b} gives ["a b"].
  • An empty dictionary returns every word as its own token.
  • A long chain "a b c d" with {a b, b c, c d} gives ["a b", "c d"], and it never gives ["a", "b c", "d"].
  • n-gram version: "san francisco bay area" with {san francisco, san francisco bay} gives ["san francisco bay", "area"], because the longest match wins.

Follow-ups

  1. Phrases of any length (true n-grams). Use leftmost-longest matching with a trie of words, or a map from first word to candidate lengths, as in ngram_split. Decide whether longest or shortest wins at a tie position, and say which one you chose.
  2. Best segmentation instead of leftmost. For plain bigrams, leftmost also maximizes the number of merges. All matches have the same length, so taking the leftmost one is the classic "earliest finish" interval-scheduling greedy. Greedy breaks once phrases have weights (for example frequency or PMI) or mixed lengths. In "a b c" with weights {a b: 1, b c: 5}, leftmost gives ["a b", "c"], but the best total weight is ["a", "b c"]. A right-to-left DP fixes this: best[i] = max(best[i+1], w(words[i:i+k]) + best[i+k]) over every phrase length k, as in best_weight_split.
  3. BPE-style ranked merges. Each bigram has a priority, and you merge the highest-priority pair anywhere in the sequence, then repeat on the merged tokens. Keep a heap of adjacent pairs keyed by rank and a doubly linked list of tokens. Mark stale heap entries invalid when their neighbors change. This runs in O(n log n). The naive rescan is O(n^2). This is a natural extension for xAI, because it is how tokenizers build subword units.
  4. Streaming input. Words arrive one at a time. Hold back at most L - 1 words, and emit a token once no longer match can still start at the head of the buffer. For bigrams, keep one pending word.
  5. Huge dictionary or many queries. Intern words to integer ids so that pair lookups hash two ints. You can also store pairs as one 64-bit key, (id1 << 32) | id2, and shard the set if it does not fit on one machine.

Interviewer expectations

Mid-level (L4)

  • Asks what happens on an overlap and whether the bigrams are directed, then writes the greedy scan with a hash set of tuples.
  • Handles a single word, the last word, and repeated words without index errors.
  • States O(n) time and walks through the overlap example.

Senior (L5)

  • Explains why the left-to-right scan is the leftmost rule, and states the invariant that every word is used exactly once and in order.
  • Extends cleanly to n-grams with leftmost-longest matching, and normalizes whitespace in both the text and the dictionary.
  • Writes tests that target overlaps and chains, not only the happy path.

Staff (L6+)

  • Proves leftmost is optimal for counting bigram merges (equal-length intervals), shows where it breaks (weights, mixed lengths), and gives the DP.
  • Connects the problem to tokenization: ranked BPE merges using a heap and a linked list, plus streaming with a bounded lookahead buffer.
  • Discusses production concerns: interning words to ids, memory for large phrase tables, and making the merge rule deterministic so training and serving produce the same tokens.

Sign in to comment, vote or remix.