CODING CHALLENGE · N°19

BPE Merge Step

Medium AI EngineeringTokenizationNLP

One step of how every tokenizer vocabulary is built: find the most frequent adjacent symbol pair across the corpus and merge it everywhere. Run this a few thousand times and you have byte-pair encoding — the algorithm behind the tokens every model reads.

The problem

Given words — a list of words, each a list of symbols — perform one BPE training step: count every adjacent symbol pair across all words, find the most frequent pair (break ties by choosing the lexicographically smallest pair), and merge every non-overlapping occurrence of that pair, left to right, into a single concatenated symbol. Return a pair [best_pair, new_words] where best_pair is the merged pair and new_words the corpus after the merge.

EXAMPLE 1
Input words = [['l','o','w'], ['l','o','w','e','r']]
Output [('l','o'), [['lo','w'], ['lo','w','e','r']]]
'lo' and 'ow' both appear twice — tie broken lexicographically
EXAMPLE 2
Input words = [['a','a','a']]
Output [('a','a'), [['aa','a']]]
non-overlapping, left to right: aa + a, not a + aa
EXAMPLE 3
Input after many steps
Output frequent substrings become single tokens
this is how vocabularies grow
CONSTRAINTS
  • Ties on count → lexicographically smallest pair wins (compare as a 2-tuple / 2-array of strings).
  • Merging is left-to-right and non-overlapping within a word.
  • Words with fewer than 2 symbols contribute no pairs and pass through unchanged.
SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement bpe_merge_step(words)[best_pair, new_words]: count adjacent pairs, pick the most frequent (lexicographically smallest on ties), merge it everywhere left-to-right without overlaps.

HINTS — 4 IDEAS
  1. Count pairs with a dict keyed by the (left, right) tuple, scanning each word once.
  2. Pick the winner by sorting on (-count, pair) or an equivalent max with a tie-break key.
  3. Merge with an index walk: if symbols i and i+1 equal the pair, emit the concatenation and jump i by 2; else emit symbols[i] and jump by 1.
  4. Check the ['a','a','a'] case: exactly one merge should happen, producing ['aa','a'].
CPython · WebAssembly
Approach, complexity & discussion — open after you solve

The approach

Count every adjacent symbol pair across the corpus, find the single most frequent pair, and merge it into one new symbol everywhere it appears — then repeat for a fixed number of merges. Each merge adds one vocabulary entry and shortens the sequences that contained that pair.

Complexity

Time O(k·n) for k merges over n corpus symbols — each merge recounts the pairs; space O(distinct pairs).

Common mistakes

  • Not recounting pair frequencies after a merge — merging changes which pairs are now adjacent.
  • Merging across word boundaries when the scheme keeps words separate.
  • Breaking frequency ties inconsistently, so the learned vocabulary is non-deterministic.

Where this shows up

This merge step is how modern tokenizers are trained — byte-pair encoding builds the vocabulary that turns “tokenization” into a couple of tokens. Understanding it demystifies why token counts, costs, and context limits behave the way they do, and why code and rare words fragment into many tokens.

Finished this one? 0 / 75 Challenges done

Explore the topic

See this alongside everything else on the same subject — handbooks, system designs, challenges and tools, in one place.

More Challenges