MMR Reranking
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.
lam = 1, k = 3indices in pure relevance ordertwo near-duplicate top docs, lam = 0.3the duplicate is skipped for a diverse dock = 1the single most query-similar index- 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.
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.
Implement mmr(q, docs, lam, k) → list of k selected indices, greedily maximizing λ·relevance − (1−λ)·redundancy with cosine similarity.
- 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.
Explore the topic
See this challenge alongside everything else on the same subject — handbooks, system designs, algorithms and tools, in one place.