11. Implement a function that selects up to k documents from candidates using Maximal Marginal Relevance with tradeoff λ ∈ [0,1] and a total token budget B. Each candidate has an embedding vector, a token cost, and an id, and you are given a query embedding; at each step pick the remaining document maximizing λ · cos(q, d) - (1-λ) · max_{s ∈ S} cos(d, s) subject to total tokens in S staying at most B.
Return selected IDs in selection order, enforce k and the token budget at every step, validate vector dimensions and λ, and apply the reported tie breaks: higher query similarity, smaller token cost, then lexicographically smaller ID.
I would use greedy Maximal Marginal Relevance selection. I first compute each document's cosine similarity to the query. Then I repeatedly consider only documents that still fit the token budget. I score each one by query relevance minus its maximum similarity to an already selected document. I break ties by higher query similarity, lower token cost, then smaller ID. After each choice, I update the redundancy values. The time is O(knD), and auxiliary space is O(n + k).
See the Code while reading this explanation.
We need to choose up to k documents. Each document has an ID, an embedding, and a token cost. We want documents that match the query, but we also want to avoid choosing documents that are too similar to ones already selected. Every chosen document must keep the total token count at or below B. At each step, we score every remaining document that still fits. We choose the best score and use fixed tie breaks. This greedy method fits because the MMR rule itself asks us to make one best choice at a time.
- Can token costs be zero? The implementation supports non-negative token costs.
- Should zero-length or zero-norm embeddings be rejected? I would reject them because cosine similarity would be undefined.
- Should the function return fewer than k documents when no remaining document fits the budget? Yes. It stops and returns the IDs already selected.
![Implement a function that selects up to k documents from candidates using Maximal Marginal Relevance with tradeoff λ ∈ [0,1] and a total token budget B. Each candidate has an embedding vector, a token cost, and an id, and you are given a query embedding; at each step pick the remaining document maximizing λ · cos(q, d) - (1-λ) · max_{s ∈ S} cos(d, s) subject to total tokens in S staying at most B. diagram](resources/diagrams/webp/cohere-ai-engineer-cohere-coding-82.webp)
The input contains the query embedding, a list of candidate documents, k, the token budget B, and λ. Each candidate has an ID, an embedding, and a token cost. The output is a list of selected IDs in selection order. The list can contain fewer than k IDs if the remaining candidates do not fit the budget.
I validate that 0 ≤ λ ≤ 1, k ≥ 0, and B ≥ 0. All embeddings must have the same non-zero dimension. The query and every document embedding must also have a non-zero L2 norm because cosine similarity divides by those norms.
I precompute sim_q[d] = cos(q, d) for every candidate. This is the relevance of each document to the query.
I start with an empty selected list, used_tokens = 0, and all candidate indices in the remaining set. I also keep max_red[d]. It means the largest cosine similarity between document d and any document already selected.
Before anything is selected, redundancy is treated as 0. After the first selection, max_red is updated from real cosine values. This matters because cosine similarity can be negative.
For each remaining document d, I first check whether used_tokens + d.tokens ≤ B. If it does not fit, I do not consider it in that step.
When S is empty, the score is: score(d) = λ · sim_q[d]
After at least one document is selected, the score is: score(d) = λ · sim_q[d] − (1 − λ) · max_red[d]
The first term rewards relevance to the query. The second term penalizes similarity to documents already selected.
I choose the feasible document with the highest MMR score. If two documents have exactly the same score, I choose the one with higher query similarity. If that also ties, I choose the smaller token cost. If that also ties, I choose the lexicographically smaller ID.
After choosing a document, I add its ID to the result, add its token cost to used_tokens, and remove it from the remaining set. Then I update max_red for every remaining document using its cosine similarity to the newly selected document.
The example uses λ = 0.7, B = 20, k = 2, and q = [1, 0].
A has embedding [1, 0] and costs 6 tokens. Its query similarity is 1.00. B has embedding [0.8, 0.6] and costs 8 tokens. Its query similarity is 0.80. C has embedding [0.8, -0.6] and costs 5 tokens. Its query similarity is 0.80.
At step 1, S is empty, so redundancy is 0. The scores are A = 0.700, B = 0.560, and C = 0.560. A has the highest score, so we choose A. used_tokens becomes 6.
At step 2, cos(B, A) = 0.800 and cos(C, A) = 0.800. Both B and C get score 0.560 − 0.240 = 0.320. Their query similarities also tie. C costs 5 tokens while B costs 8, so the token-cost tie break chooses C. used_tokens becomes 11.
Now |S| = k = 2, so we stop. The returned IDs are ["A", "C"].
At every iteration, the feasible set contains exactly the remaining documents whose addition would keep the total token count at or below B. For each of those documents, max_red stores its true maximum cosine similarity to the selected set. Therefore, the computed score is exactly the MMR score required by the question.
The algorithm picks the feasible document with the largest score and applies the required tie breaks in the required order. The chosen document is then removed, so it cannot be selected twice. The loop stops when k documents have been selected or no remaining document fits the budget.
Let n be the number of candidates and D be the embedding dimension. Computing query similarities costs O(nD). Each of up to k selections can compare the newly selected embedding with up to n remaining embeddings, which costs O(knD). So total time is O(nD + knD) = O(knD). If k = n, the worst case is O(n²D).
The main extra structures store query similarities, maximum redundancy values, remaining indices, and selected indices. Auxiliary space is O(n + k), excluding the input embeddings and returned IDs.
Important edge cases are k = 0, an empty candidate list, no document fitting the remaining budget, zero token-cost documents, invalid λ, mismatched embedding dimensions, and zero-norm embeddings.
The key idea is to keep one relevance value and one redundancy value for each candidate. Query similarity is computed once. Redundancy is updated incrementally after every selection, so we do not need to store an n × n similarity matrix. The central invariant is: before each selection, max_red[d] equals the maximum cosine similarity between remaining document d and every document already selected. Because the algorithm also filters by the remaining token budget before scoring, every chosen document keeps the total cost within B. This directly implements the greedy MMR rule shown in the diagram.
from __future__ import annotations
from dataclasses import dataclass
import math
from typing import Sequence
@dataclass(frozen=True)
class Document:
id: str
embedding: Sequence[float]
tokens: int
def mmr_select(
candidates: Sequence[Document],
query: Sequence[float],
k: int,
B: int,
lam: float,
) -> list[str]:
"""Return selected document IDs in MMR selection order."""
# Validate scalar inputs before doing vector calculations.
if not 0.0 <= lam <= 1.0:
raise ValueError("lam must be in [0, 1]")
if k < 0:
raise ValueError("k must be >= 0")
if B < 0:
raise ValueError("B must be >= 0")
n = len(candidates)
if n == 0 or k == 0:
return []
# Cosine similarity needs a non-empty query vector.
dimension = len(query)
if dimension == 0:
raise ValueError("query embedding is empty")
# Reject a zero-norm query because cosine similarity would divide by zero.
query_norm = math.sqrt(sum(x * x for x in query))
if query_norm == 0.0:
raise ValueError("query embedding has zero norm")
# Check dimensions, token costs, and document norms before selection starts.
document_norms: list[float] = []
for doc in candidates:
if len(doc.embedding) != dimension:
raise ValueError("all embeddings must have the same dimension")
if doc.tokens < 0:
raise ValueError("token cost must be non-negative")
norm = math.sqrt(sum(x * x for x in doc.embedding))
if norm == 0.0:
raise ValueError(f"embedding for doc {doc.id} has zero norm")
document_norms.append(norm)
def dot(a: Sequence[float], b: Sequence[float]) -> float:
# Dimensions were validated above, so paired elements are safe to use.
return sum(x * y for x, y in zip(a, b))
def cosine_between_indices(i: int, j: int) -> float:
# Reuse stored document norms during incremental redundancy updates.
return dot(candidates[i].embedding, candidates[j].embedding) / (
document_norms[i] * document_norms[j]
)
# Query relevance never changes, so calculate it once per candidate.
sim_q: list[float] = [
dot(query, doc.embedding) / (query_norm * document_norms[i])
for i, doc in enumerate(candidates)
]
# max_red[i] will become the maximum similarity from candidate i to S.
# Start at -inf so a negative first cosine value is stored correctly.
max_red: list[float] = [float("-inf")] * n
selected_indices: list[int] = []
selected_ids: list[str] = []
remaining: set[int] = set(range(n))
used_tokens = 0
# Greedily choose at most k documents.
while len(selected_indices) < k and remaining:
# Only documents that keep the total token count within B are feasible.
feasible = [i for i in remaining if used_tokens + candidates[i].tokens <= B]
# Stop when no remaining document can fit the token budget.
if not feasible:
break
best_idx: int | None = None
best_score = float("-inf")
best_sim = float("-inf")
best_tokens = math.inf
best_id = ""
# Score every feasible document using the current selected set S.
for i in feasible:
doc = candidates[i]
# The MMR definition uses zero redundancy when S is empty.
redundancy = 0.0 if not selected_indices else max_red[i]
score = lam * sim_q[i] - (1.0 - lam) * redundancy
# Choose the highest score. For exact score ties, apply the required
# order: higher query similarity, smaller token cost, smaller ID.
is_better = (
best_idx is None
or score > best_score
or (
score == best_score
and (
sim_q[i] > best_sim
or (
sim_q[i] == best_sim
and (
doc.tokens < best_tokens
or (doc.tokens == best_tokens and doc.id < best_id)
)
)
)
)
)
# Save all comparison values for the current best document.
if is_better:
best_idx = i
best_score = score
best_sim = sim_q[i]
best_tokens = doc.tokens
best_id = doc.id
# The feasible list is non-empty, so this is only a defensive guard.
if best_idx is None:
break
# Commit the choice and remove it so the same document cannot repeat.
selected_indices.append(best_idx)
selected_ids.append(candidates[best_idx].id)
remaining.remove(best_idx)
used_tokens += candidates[best_idx].tokens
# Update each remaining document's maximum similarity to selected docs.
for j in remaining:
similarity_to_new = cosine_between_indices(j, best_idx)
# The first selected document establishes the first real redundancy.
if len(selected_indices) == 1:
max_red[j] = similarity_to_new
else:
# Later selections keep the largest similarity seen so far.
max_red[j] = max(max_red[j], similarity_to_new)
# Return IDs in the exact order in which the documents were selected.
return selected_ids
def main() -> None:
# Run the same verified example used in the diagram.
candidates = [
Document(id="A", embedding=[1.0, 0.0], tokens=6),
Document(id="B", embedding=[0.8, 0.6], tokens=8),
Document(id="C", embedding=[0.8, -0.6], tokens=5),
]
query = [1.0, 0.0]
result = mmr_select(
candidates=candidates,
query=query,
k=2,
B=20,
lam=0.7,
)
print(result) # ['A', 'C']
if __name__ == "__main__":
main()Let n be the number of candidates and D be the embedding dimension. Computing every document's similarity to the query takes O(nD). We can make at most k selections. After each selection, we may compare the chosen document with up to n remaining documents, and each cosine calculation touches D embedding values. That gives O(knD) work. Including preprocessing, the total is O(nD + knD) = O(knD). If k = n, this becomes O(n²D). Auxiliary space is O(n + k), excluding the input embeddings and returned IDs.
This pattern is useful when a retrieval system has many relevant documents but should return a smaller set that is both relevant and varied. For example, after semantic retrieval, MMR can choose context documents for an AI application while avoiding near-duplicate passages and staying inside a token budget.
This problem tests whether you can turn a mathematical ranking rule into correct code. The interviewer can check if you understand cosine similarity, greedy selection, state updates, deterministic tie breaking, and a hard token budget. It also tests whether you notice less obvious cases such as negative cosine values, zero-norm vectors, and zero token costs. Finally, you must explain why the incremental redundancy state is correct and give time and space complexity that matches the actual implementation.
- Forgetting to filter candidates by the remaining token budget before choosing the best score. That can make the final selection exceed B.
- Using only query similarity. MMR must also subtract the maximum similarity to the already selected documents.
- Initializing later redundancy updates as if cosine similarity could not be negative. Cosine similarity can be negative, so the first real redundancy value must not be incorrectly clamped to zero.
- Applying the tie breaks in the wrong order. The required order is higher query similarity, then smaller token cost, then lexicographically smaller ID.
- Forgetting to remove the chosen candidate from the remaining set. That can select the same document again.
- Claiming O(n²) extra memory. The incremental max_red design uses O(n + k) auxiliary space and does not store a full pairwise similarity matrix.
Explain the state before writing code. Say that sim_q never changes, while max_red changes after each selection. Then state the invariant: max_red[d] is the largest similarity between remaining document d and any selected document. This makes the scoring rule, state update, and O(n + k) auxiliary-space claim easy to justify.