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

Explore the topic

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