CODING CHALLENGE · N°20

MMR Reranking

Medium AI EngineeringRetrievalRAG

Top-k by similarity returns five copies of the same paragraph. Maximal Marginal Relevance fixes it: pick results that are relevant to the query AND different from what you already picked. The greedy diversity algorithm inside serious RAG pipelines.

The problem

Given a query vector q, candidate document vectors docs, a tradeoff lam (λ ∈ [0,1]) and k, return the indices of the k documents selected by Maximal Marginal Relevance: repeatedly pick the unselected document maximizing lam · cos(q, d) − (1 − lam) · max(cos(d, s) for s in selected), where the penalty term is 0 for the first pick. Use cosine similarity. Break score ties by the lower index.

EXAMPLE 1

Input lam = 1, k = 3

Output indices in pure relevance order

λ=1 disables the diversity penalty

EXAMPLE 2

Input two near-duplicate top docs, lam = 0.3

Output the duplicate is skipped for a diverse doc

diversity outweighs raw relevance

EXAMPLE 3

Input k = 1

Output the single most query-similar index

first pick is always pure relevance

CONSTRAINTS

Return document indices (0-based), in selection order.
First selection = highest cos(q, d) — no penalty exists yet.
Tie on MMR score → the smaller index wins.
k ≤ len(docs); vectors are non-zero.

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 mmr(q, docs, lam, k) → list of k selected indices, greedily maximizing λ·relevance − (1−λ)·redundancy with cosine similarity.

HINTS — 4 IDEAS

Write a cosine helper first: dot(a,b) / (|a|·|b|).
Keep a selected list; each round score every unselected candidate and take the argmax.
The redundancy term is the MAX similarity to any already-selected doc — not the average.
λ=1 should reduce to plain top-k by query similarity; use that as your sanity check.

CPython · WebAssembly

Approach, complexity & discussion — open after you solve

The approach

Build the result set greedily, balancing relevance against novelty. At each step pick the candidate that maximizes λ·relevance(query) − (1−λ)·max similarity to already-selected. The first term pulls in relevant items; the second penalizes anything too similar to what you already have. λ tunes the relevance-vs-diversity trade-off.

Complexity

Time O(k·n) to select k items from n candidates — each pick scans the remaining candidates and compares against the selected set.

Common mistakes

Ranking on relevance alone — you get five near-duplicate passages and waste the budget.
Setting λ wrong: all relevance is redundant, all diversity drifts off-topic.
Computing the penalty against the query instead of against the already-selected items.

Where this shows up

MMR diversifies RAG context and search results so the prompt is not stuffed with paraphrases of one passage — you cover more of the answer within the same token budget. It is a standard reranking step in retrieval and in extractive summarization.

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.

RAG & Retrieval