0

Greedy longest-match tokenizer

by Frontier Interview· 1 day ago· ⑂ 0 remixes

Question: Greedy longest-match tokenizer

Companies: Anthropic, Google

Difficulty: Medium

Problem statement

Implement a greedy longest-match tokenizer.

You are given a string text and a list dictionary. Each entry has the form "<key>:<id>". key is a non-empty token string and id is its identifier. The key may itself contain :, but the id never does, so the id is everything after the last colon.

Write tokenize(text: str, dictionary: List[str]) -> List[str] that scans text from left to right and returns the ids of the tokens it produces:

  • At the current position, find every key that starts there and choose the longest one. Emit its id and move past it.
  • If no key starts at the current position, emit "[UNK]" for that single character and move forward by one.
  • If the same key appears more than once in dictionary, use the first id.
  • Raise an error for a malformed entry (no colon, empty key or empty id).

Matching is case-sensitive. text may be empty.

Constraints: 0 <= len(text) <= 10^5, up to 10^4 dictionary entries, each key at most 50 characters.

Example 1

text = "catalog"
dictionary = ["cat:10", "catalog:11", "a:12", "log:13"]
output = ["11"]

catalog is longer than cat, so it wins at index 0.

Example 2

text = "cats!"
dictionary = ["cat:10", "s:14", "c:15"]
output = ["10", "14", "[UNK]"]

cat beats c, then s matches, and ! has no key.

Aim for better than trying every substring length with slicing at each position. State the time complexity of your solution.

Approach

The vocabulary arrives as strings of the form "<key>:<id>". Parse each one by splitting on the last colon (rsplit(":", 1)), because a key may itself contain a colon ("a:b:7" is key a:b, id 7) but an id never does. If a key appears twice, keep the first id. Ask about this before you code it.

Then tokenize greedily from left to right. At position i, find the longest key that starts at i, emit its id, and jump past it. If no key starts at i, emit the unknown marker [UNK] for that one character and advance by one. The scan always moves forward, so it terminates, and every character of text is covered by exactly one output token.

Brute force. Put the keys in a hash map and remember the longest key length L. At each i, try text[i:i+L], then text[i:i+L-1], and so on down to length 1. Stop at the first hit. This is correct, but every slice copies up to L characters, so one position costs O(L^2) and the whole pass costs O(n * L^2). Say this first. It is a fine answer when the keys are short.

Trie. Build a trie over the keys and store an id on every node where a key ends. At position i, walk down the trie along text[i], text[i+1], ... until the next character has no child, and remember the deepest node you passed that ends a key. That node is the longest match. The walk never goes further than L characters and never copies a string, so one position costs O(L) and the pass costs O(n * L).

Why greedy is the spec, not an optimization. Longest-match-first does not minimize the number of tokens. With keys ab, abc and cd, the text abcd greedily becomes abc, [UNK], but ab, cd would avoid the unknown. The problem asks for greedy matching (as WordPiece-style tokenizers do), so implement greedy. Mention the difference, because the interviewer often asks for the optimal split next (see follow-ups).

Invariants.

  • i only increases, by the length of the match or by 1.
  • The concatenation of the matched keys and the [UNK] characters, in order, rebuilds text exactly.
  • The match at i is the longest key that is a prefix of text[i:]. The trie walk only stops when no longer key can share the prefix.

Code

from typing import Dict, List, Optional, Tuple

UNK = "[UNK]"


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

    def __init__(self) -> None:
        self.children: Dict[str, "TrieNode"] = {}
        self.token_id: Optional[str] = None  # set when a key ends here


def parse_entry(entry: str) -> Tuple[str, str]:
    """Split "<key>:<id>" on the last colon, so keys may contain ':'."""
    key, sep, token_id = entry.rpartition(":")
    if not sep or not key or not token_id:
        raise ValueError(f"bad dictionary entry: {entry!r}")
    return key, token_id


class Tokenizer:
    def __init__(self, dictionary: List[str]) -> None:
        self.root = TrieNode()
        for entry in dictionary:
            key, token_id = parse_entry(entry)
            node = self.root
            for ch in key:
                node = node.children.setdefault(ch, TrieNode())
            if node.token_id is None:  # first definition of a key wins
                node.token_id = token_id

    def longest_match(self, text: str, start: int) -> Tuple[int, Optional[str]]:
        """Return (length, id) of the longest key at text[start:], or (0, None)."""
        node = self.root
        best_len, best_id = 0, None
        i = start
        while i < len(text):
            node = node.children.get(text[i])
            if node is None:
                break
            i += 1
            if node.token_id is not None:
                best_len, best_id = i - start, node.token_id
        return best_len, best_id

    def tokenize(self, text: str) -> List[str]:
        out: List[str] = []
        i = 0
        while i < len(text):
            length, token_id = self.longest_match(text, i)
            if token_id is None:
                out.append(UNK)  # no key starts here: skip one character
                i += 1
            else:
                out.append(token_id)
                i += length
        return out


def tokenize(text: str, dictionary: List[str]) -> List[str]:
    return Tokenizer(dictionary).tokenize(text)


if __name__ == "__main__":
    vocab = ["a:1", "ab:2", "abc:3", "c:4", "x:y:9"]
    assert tokenize("abcab", vocab) == ["3", "2"]
    assert tokenize("abd", vocab) == ["2", UNK]
    assert tokenize("", vocab) == []
    assert tokenize("x:y", vocab) == ["9"]
    assert tokenize("x:", vocab) == [UNK, UNK]
    assert tokenize("aa", ["a:1", "a:7"]) == ["1", "1"]
    assert tokenize("abcd", ["ab:1", "abc:2", "cd:3"]) == ["2", UNK]
    print("ok")

Complexity

Let n = len(text), L the longest key length, and S the total length of all keys.

  • Build: O(S) time and O(S) space for the trie, at worst one node per key character.
  • Tokenize: O(n * L) time. Each position walks at most L trie edges, and only positions where a token starts are visited, so it is often closer to O(n). The output takes O(n) extra space.
  • Hash-map brute force, for comparison: O(n * L^2) time because of slicing, O(S) space.

Test cases

  • text = "abcab", dictionary ["a:1", "ab:2", "abc:3", "c:4"] gives ["3", "2"]. abc beats ab and a at index 0.
  • text = "abd", same dictionary, gives ["2", "[UNK]"]. d has no key.
  • text = "" gives [].
  • text = "zzz", any dictionary without z, gives ["[UNK]", "[UNK]", "[UNK]"].
  • text = "x:y", dictionary ["x:y:9"], gives ["9"]. The key contains a colon.
  • text = "aa", dictionary ["a:1", "a:7"], gives ["1", "1"]. The first definition wins.
  • text = "abcd", dictionary ["ab:1", "abc:2", "cd:3"], gives ["2", "[UNK]"]. Greedy is not globally optimal.
  • text = "aaaa", dictionary ["aaa:1", "a:2"], gives ["1", "2"].
  • An entry with no colon, such as "abc", raises ValueError.

Follow-ups

  1. Detokenize. Given ids, rebuild the text. Keep a reverse map from id to key. A plain [UNK] loses the original character, so change the output to carry it, for example ("[UNK]", "d"), or store the raw byte. Then detokenize(tokenize(t)) == t holds. Coalescing a run of unknown characters into one token is a common variant. Keep a buffer, and flush it when a match is found.
  2. Minimum number of tokens instead of greedy. Use DP over positions: best[i] is the fewest tokens covering text[:i]. From each i, walk the trie and relax best[i + len] for every key that ends along the walk. Unknown characters cost one token, or a large penalty if you want to avoid them. The cost is O(n * L) time and O(n) space. Recover the split with parent pointers.
  3. Streaming input. The text arrives in chunks. A match can cross a chunk boundary, so you cannot emit a token until you know that no longer key can follow. Keep the unconsumed tail (at most L - 1 characters) and the trie state. Emit once the walk falls off the trie, and flush at end of stream.
  4. Very large vocabularies or many tokenizers. Compress the trie (radix tree or double-array trie) to save memory. Or build an Aho-Corasick automaton, which lets one left-to-right pass find matches without restarting the walk at every position. This is how production longest-match tokenizers stay linear.
  5. BPE instead of a fixed dictionary. The vocabulary is a ranked list of merges. Start from single characters (or bytes) and repeatedly merge the adjacent pair with the best rank. Use a heap of candidate pairs and a linked list of symbols, which gives O(n log n) per word.

Interviewer expectations

Mid-level (L4)

  • Parses the entries correctly and asks about duplicate keys and characters that match nothing.
  • Writes the hash-map version that tries lengths from longest to shortest, with correct index handling and no infinite loop.
  • Tests the empty string, unknown characters, and a case where a shorter key is a prefix of a longer one.

Senior (L5)

  • Moves to a trie unprompted and states O(n * L) time against O(n * L^2) for slicing.
  • Handles keys that contain the delimiter (splits on the last colon) and malformed entries.
  • Explains with a counterexample why greedy differs from the minimum-token split, and sketches the DP.
  • Makes detokenization lossless.

Staff (L6+)

  • Discusses streaming across chunk boundaries, memory for a large vocabulary, and Aho-Corasick or a double-array trie.
  • Connects the problem to real tokenizers (WordPiece greedy matching, BPE merges, byte fallback for unknown text) and their tradeoffs.
  • Defines the API contract: which unknown policy to use, whether output is ids or (id, span) pairs, Unicode normalization, and thread-safe reuse of a built tokenizer.

Sign in to comment, vote or remix.