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

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

CPython · WebAssembly

Approach, complexity & discussion — open after you solve

The approach

Maintain the top b partial sequences (“beams”) ranked by score. Each step, expand every beam by every possible next token, score all the resulting candidates, and keep the best b across the whole pool. Stop at a length limit or an end token, and return the highest-scoring complete sequence. Scores are summed log-probabilities, not raw probabilities.

Complexity

Time O(steps × b × V) with vocabulary V; space O(b × length) for the kept beams.

Common mistakes

Keeping the top-b per beam instead of the top-b across all expansions — that is the crux of beam search.
Multiplying raw probabilities instead of summing log-probs — you underflow and rank wrongly.
Not length-normalizing the score, which biases the search toward shorter sequences.

Where this shows up

Beam search is a core decoding strategy for machine translation, speech recognition, and constrained generation — the middle ground between greedy decoding (b = 1) and infeasible exhaustive search. The “keep the top-b hypotheses and prune the rest” idea recurs in any large pruned search problem.

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.

LLM Engineering