Vibe Engines Learn
CODING CHALLENGE · N°24

Beam Search Decoder

Medium AI EngineeringDecodingSearch

Greedy decoding takes the best next token and never looks back — and walks straight into garden paths. Beam search keeps the k best partial sequences alive at every step. Implement the decoder and watch k=2 escape a trap that k=1 falls into.

The problem

Implement beam search over a Markov token model: given start (the initial token id), trans — a matrix where trans[i][j] is the probability of token j following token i — a number of steps, and a beam width k, return the most probable sequence of length steps + 1 (including the start token). A sequence’s score is the product of its transition probabilities. At each step, extend every beam with every possible next token, then keep only the k highest-scoring sequences (break score ties by preferring the lexicographically smaller sequence). Return the best sequence in the final beam.

EXAMPLE 1
Input start = 0, trans = [[0, 0.6, 0.4], [0, 0.5, 0.5], [0, 1.0, 0]], steps = 2, k = 1
Output [0, 1, 1] or [0, 1, 2] — score 0.30
greedy grabs the 0.6 branch and caps out at 0.30
EXAMPLE 2
Input same, k = 2
Output [0, 2, 1] — score 0.40
the beam keeps the 0.4 branch alive and wins
EXAMPLE 3
Input k ≥ number of tokens
Output exhaustive search — the true optimum
width buys optimality, at width cost
CONSTRAINTS
  • Score = product of transition probabilities along the sequence.
  • Keep exactly min(k, candidates) beams per step, highest score first.
  • Score ties → the lexicographically smaller sequence wins.
  • trans rows may contain zeros; zero-probability extensions are legal but score 0.
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 beam_search(start, trans, steps, k) → the highest-scoring token sequence, expanding all beams each step and keeping the top k.

HINTS — 4 IDEAS
  1. Represent each beam as (sequence, score); start with one beam: ([start], 1.0).
  2. Each step: for every beam and every next token j, extend to score · trans[last][j].
  3. Sort candidates by (−score, sequence) and slice the top k — that tuple sort handles the tie-break.
  4. k = 1 must reproduce greedy decoding; test against that mentally first.
CPython · WebAssembly

Explore the topic

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

LLM Engineering