Maximal Marginal Relevance Selection

mediumPROContext & Agent MemoryAI AgentsAI Agents 1

Problem Statement

Select a diverse yet relevant set of memories with Maximal Marginal Relevance (MMR): greedily pick items that are relevant to the query but not too similar to what you already picked.

Background

MMR selects the next item maximizing

$\\text{MMR} = \\lambda \\cdot rel(i) - (1 - \\lambda) \\max_{j \\in S} sim(i, j)$

where S is the already-selected set. The first pick is the highest-relevance item (the max-sim term is 0 when S is empty). Ties are broken by smaller index.

Your Task

Implement:

def mmr(relevance, similarity, lam, k):

relevance: list of floats, relevance[i] = rel of item i.
similarity: n x n nested list, similarity[i][j].
Return the list of selected indices in selection order, length min(k, n).

Input Format

relevance (list), similarity (nested list), lam (float), k (int).

Output Format

A list of ints (selected indices).

Sample

rel = [0.9, 0.8, 0.7]
sim = [[1,0.9,0.1],[0.9,1,0.2],[0.1,0.2,1]]
print(mmr(rel, sim, 0.5, 2))

Output:

[0, 2]

Example

Input

rel = [0.9, 0.8, 0.7]
sim = [[1,0.9,0.1],[0.9,1,0.2],[0.1,0.2,1]]
print(mmr(rel, sim, 0.5, 2))

Output

[0, 2]

Reasoning

First Selection (Empty Set): Since the selected set $S$ is empty, the similarity penalty is 0 for all items. We calculate the MMR score as $\lambda \cdot rel(i)$ with $\lambda = 0.5$ :
- Item 0: $0.5 \cdot 0.9 = 0.45$
- Item 1: $0.5 \cdot 0.8 = 0.40$
- Item 2: $0.5 \cdot 0.7 = 0.35$
- Item 0 has the highest score, so it is selected first. $S = [0]$ .
Second Selection (S = {0}): We evaluate the remaining items (1 and 2) using the formula $\text{MMR} = 0.5 \cdot rel(i) - 0.5 \cdot \max_{j \in S} sim(i, j)$ . We must find the maximum similarity between each candidate and the already selected item 0:
- Item 1: $sim(1, 0) = 0.9$ . Score: $0.5 \cdot 0.8 - 0.5 \cdot 0.9 = 0.4 - 0.45 = -0.05$ .
- Item 2: $sim(2, 0) = 0.1$ . Score: $0.5 \cdot 0.7 - 0.5 \cdot 0.1 = 0.35 - 0.05 = 0.30$ .
Comparison and Selection: Comparing the scores from the second step, Item 2 ( $0.30$ ) is greater than Item 1 ( $-0.05$ ). Therefore, Item 2 is selected next. $S = [0, 2]$ .
Termination: The target size $k=2$ is reached, so the algorithm stops.
The final output is [0, 2]

Constraints

0 <= lam <= 1, similarity is n x n.
First pick = argmax relevance (ties: smaller index).
Each later pick maximizes lam*rel - (1-lam)*max sim to selected; ties: smaller index.

Learn this in

Memory II: Hybrid Stores, Mem0 & Skill LibrariesAgent Engineering Study Plan

More from AI Agents 1

Cheapest Cache TTL Meeting a Freshness SLAhard Bootstrap Confidence Interval Width for Eval Accuracymedium Route to the Cheapest Model Meeting Qualitymedium pass@k Estimatormedium Savings from a Semantic Cachemedium

Concept primerPREVIEW

1. Background Knowledge

Maximal Marginal Relevance (MMR) is a classic diversification technique used in information retrieval and, increasingly, in context & agent memory systems. When an AI agent retrieves memories or documents for a query, simply picking the top- $k$ most relevant items often yields redundant results (e.g., three nearly identical paragraphs). MMR balances relevance (how well an item matches the query) against diversity (how different an item is from items already chosen).

The objective function is:

\text{MMR}(i) = \lambda \cdot \text{rel}(i) - (1 - \lambda) \cdot \max_{j \in S} \text{sim}(i, j)

where $S$ is the set of already-selected indices. The parameter $\lambda \in [0, 1]$ controls the trade-off: $\lambda = 1$ ignores diversity (pure relevance), while $\lambda = 0$ ignores relevance (pure diversity). When $S$ is empty, the max-sim term is defined as $0$ , so the first pick is simply the item with the highest relevance.

In the context of agent memory, MMR ensures that the retrieved context window contains varied, complementary pieces of information rather than redundant duplicates, improving the quality of downstream reasoning.

2. Algorithm Approach

This is a greedy selection algorithm. At each step, you evaluate every unselected candidate item, compute its MMR score given the current selected set $S$ , and pick the candidate with the highest score. Ties are broken by choosing the smaller index. You repeat this until you have selected $\min(k, n)$ items.

The key insight is that the "max similarity to already-selected items" term requires, for each candidate $i$ , scanning all previously selected items $j$ and taking the maximum of $\text{sim}(i, j)$ . This makes the per-candidate cost proportional to $|S|$ .

3. Step-by-Step Strategy

Initialize: Create an empty list selected and a boolean array used of length $n$ (all False).
Loop for t from $0$ to $\min(k, n) - 1$ :
For each candidate index $i$ where used[i] is False:
Compute rel_score = lam * relevance[i].
If selected is non-empty, compute max_sim = max(similarity[i][j] for j in selected). Otherwise, max_sim = 0.
Compute mmr_score = rel_score - (1 - lam) * max_sim.
Among all candidates, find the one with the highest mmr_score. If there is a tie, choose the smallest index.
Append that index to selected and mark used[i] = True.
Return the selected list.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results