Apply ordered token merges to a string

This is the merge step of byte-pair encoding (BPE). The alphabet is a merge list: entry i says "wherever two adjacent tokens spell alphabet[i], replace them with i". Merges apply in alphabet order, one full pass each.

Pin down the rules first. Ask the interviewer, then say them back:

• The single-character prefix gives the start sequence, one id per character. A character with no token is an error.
• For each later token t (id i), scan the current sequence left to right. Where two adjacent ids x, y have alphabet[x] + alphabet[y] == t, replace the pair with i and continue after it. Matches do not overlap: "aaa" with merge "aa" gives [aa, a], not [a, aa].
• t may split in more than one place ("abc" = a + bc = ab + c). Any split into two earlier tokens counts.
• A token with no split into two earlier tokens can never match. It is skipped. You do not need to check this separately: the sequence only ever holds earlier ids, so such a pass finds nothing.

Brute force. One left-to-right pass per merge token, rebuilding the list. With N characters and M merge tokens that is O(M * N) pair checks (each check compares strings of up to L characters). It is correct and a good first answer; write it, then use it as the oracle for the faster version.

Optimal: linked list plus a pair index. Most merges touch a tiny part of the text, so do not rescan it.

1. Keep the sequence as a doubly linked list over the original character positions. A node keeps the position of its first character, so position order is text order forever.
2. Keep where[(x, y)]: the set of node positions p where node p has id x and its next node has id y.
3. For merge token i, list its splits (a, b) with both parts earlier tokens. Pop every position for those pairs, sort them, and process them in order.
4. For each position, re-check that the node is alive and that it and its next node still spell t. A position can be stale because an earlier merge in this same pass consumed one of its nodes (the overlap case). Skip stale entries.
5. Merge: the left node takes id i, the right node is unlinked. Remove the two neighbouring pairs that disappear and add the two new ones (prev, i) and (i, next).

Why it is correct. Processing sorted candidates and skipping stale ones is exactly the brute-force left-to-right scan. A merge in this pass only creates pairs that contain i, and i can never be a part of its own split, so no new candidate appears mid-pass.

Invariants. ids along the list is always the current sequence. Every live adjacent pair is in where (stale extra entries are allowed and filtered on use). The list only shrinks, so there are at most N - 1 merges in total.

from collections import defaultdict
from typing import Dict, List, Set, Tuple


def _base_ids(text: str, alphabet: List[str]) -> Tuple[Dict[str, int], List[int]]:
    char_id: Dict[str, int] = {}
    k = 0
    while k < len(alphabet) and len(alphabet[k]) == 1:
        char_id[alphabet[k]] = k
        k += 1
    try:
        return char_id, [char_id[c] for c in text]
    except KeyError as e:
        raise ValueError(f"character {e} has no single-character token") from None


def tokenize_brute(text: str, alphabet: List[str]) -> List[int]:
    """O(M * N): one left-to-right pass over the sequence per merge token."""
    char_id, seq = _base_ids(text, alphabet)
    for tid in range(len(char_id), len(alphabet)):
        target, out, j = alphabet[tid], [], 0
        while j < len(seq):
            if j + 1 < len(seq) and alphabet[seq[j]] + alphabet[seq[j + 1]] == target:
                out.append(tid)
                j += 2
            else:
                out.append(seq[j])
                j += 1
        seq = out
    return seq


def tokenize(text: str, alphabet: List[str]) -> List[int]:
    """Linked list over positions + an index from adjacent pair to positions."""
    char_id, ids = _base_ids(text, alphabet)
    n = len(ids)
    nxt = list(range(1, n + 1))          # n means "no next node"
    prv = list(range(-1, n - 1))
    alive = [True] * n
    where: Dict[Tuple[int, int], Set[int]] = defaultdict(set)
    for p in range(n - 1):
        where[(ids[p], ids[p + 1])].add(p)

    first_id = {s: i for i, s in reversed(list(enumerate(alphabet)))}

    for tid in range(len(char_id), len(alphabet)):
        target = alphabet[tid]
        # Every way to write target as (earlier token, earlier token).
        pairs = []
        for k in range(1, len(target)):
            a, b = first_id.get(target[:k]), first_id.get(target[k:])
            if a is not None and b is not None and a < tid and b < tid:
                pairs.append((a, b))
        starts: List[int] = []
        for pr in pairs:
            starts.extend(where.pop(pr, ()))
        # Position index = original char offset, so sorting restores text order.
        for p in sorted(starts):
            q = nxt[p] if alive[p] else n
            if q == n or alphabet[ids[p]] + alphabet[ids[q]] != target:
                continue  # stale: p or its neighbour was merged earlier in this pass
            left, right = prv[p], nxt[q]
            if left >= 0:
                where[(ids[left], ids[p])].discard(left)
            if right < n:
                where[(ids[q], ids[right])].discard(q)
            ids[p] = tid
            alive[q] = False
            nxt[p] = right
            if right < n:
                prv[right] = p
                where[(tid, ids[right])].add(p)
            if left >= 0:
                where[(ids[left], tid)].add(left)

    out, p = [], 0
    while p < n:
        out.append(ids[p])
        p = nxt[p]
    return out


if __name__ == "__main__":
    A = ["a", "b", "c", "ab", "bc", "abc", "cab"]
    assert tokenize("abcab", A) == tokenize_brute("abcab", A) == [5, 3]
    assert tokenize("aaa", ["a", "aa"]) == [1, 0]          # no overlap
    assert tokenize("aaaa", ["a", "aa", "aaaa"]) == [2]
    assert tokenize("ab", ["a", "b", "xyz", "ab"]) == [3]  # unusable token skipped
    assert tokenize("abc", ["a", "b", "c", "bc", "abc", "ab"]) == [4]
    assert tokenize("", ["a"]) == []

Let N = len(text), M = number of multi-character tokens, L = longest token length.

• Brute force: O(M * N * L) time, O(N) extra space.
• Optimal, setup: O(N) to map characters and build the pair index; O(total alphabet length) for the string-to-id map.
• Optimal, per merge token: O(L^2) to build its splits (slicing and hashing each part), plus O(c log c) for its c candidate positions. Each merge then does O(1) index updates.
• Total: candidates come from pairs that existed, and the sequence loses one node per merge, so the number of real merges is at most N - 1. Stale entries are at most a constant factor more. Time is O(M * L^2 + N log N); space is O(N) for the list and index.

• text = "abcab", alphabet [a, b, c, ab, bc, abc, cab] -> [5, 3] (ab gives [ab, c, ab]; bc finds nothing; abc matches ab + c; cab finds nothing).
• text = "abc", alphabet [a, b, c, bc, abc, ab] -> [4] (bc merges first, then abc matches a + bc; ab comes too late).
• text = "aaa", alphabet [a, aa] -> [1, 0]: left to right, no overlap.
• text = "aaaa", alphabet [a, aa, aaaa] -> [2]: aa gives [aa, aa], then aaaa = aa + aa.
• text = "ab", alphabet [a, b, xyz, ab] -> [3]: xyz has no split and is skipped.
• text = "" -> [].
• text = "ad", alphabet [a, b] -> ValueError (no token for d).
• Only single-character tokens -> the start sequence unchanged.
• Random check: 3,000 random small alphabets and texts give the same output from tokenize and tokenize_brute.

• Decode. Map ids back to strings and join. Then ask the candidate to show decode(tokenize(s)) == s: every merge keeps the concatenated text the same.
• Real BPE priority. Instead of one pass per merge in order, repeatedly merge the adjacent pair with the lowest merge rank anywhere in the sequence. Use a min-heap of (rank, position) over live pairs, with the same linked list and lazy deletion of stale heap entries. O(N log N) per text, independent of M.
• Train the merges. Given a corpus, learn the alphabet: count adjacent pairs, merge the most frequent, update counts incrementally with the same pair index (plus a heap or bucket queue keyed by count). This is the original BPE training loop.
• Streaming or many texts. Tokenize text that arrives in chunks: keep a tail that may still merge with the next chunk (bounded by L characters), and emit the rest. For many short texts, cache results per pre-split word, as production tokenizers do.
• Bytes instead of characters. Start from the 256 byte values of the UTF-8 encoding so any input is valid and the unknown-character error goes away.

Mid-level (L4)

• Asks how overlapping matches and multiple splits are handled, and gets the left-to-right rule right.
• Writes a correct brute force with clean pass logic and tests the "aaa" case.
• States O(M * N) and sees that it rescans text that cannot change.

Senior (L5)

• Moves to the linked list plus pair index without prompting, and keeps it correct under stale entries.
• Explains why no new candidate for the current token can appear mid-pass.
• Uses the brute force as an oracle in a randomized test.
• Names this as BPE and sketches the rank-based heap version.

Staff (L6+)

• Discusses how real tokenizers differ: rank priority, byte-level fallback, pre-tokenization into words, and per-word caching for throughput.
• Reasons about reversibility, determinism across versions of the merge list, and what breaks if the alphabet changes after a model is trained on it.
• Covers training the merge list and the cost of incremental pair-count updates on a large corpus.