Citation Highlighting

Problem Overview

You are given a document produced by an LLM and a list of source phrases. Your task is to find exact phrase matches in the document and wrap the matched regions in <yellow> and </yellow> tags.

This is usually asked in two parts:

1. First, find all matching phrases and merge overlapping or adjacent matches into highlighted regions.
2. Then, count how often each source appears across the entire document and add source citations to every highlighted region.

The core problem is interval merging. The follow-up tests whether you can preserve the relationship between merged intervals and the source phrases that contributed to them.

The core requirements are stable: exact word-level matching, merging overlaps, global source counts, and citations for merged regions. Details such as case sensitivity, punctuation rules, how adjacency is defined, and how frequency ties are ordered vary between interviewers, so confirm them before coding. The choices used by this write-up are stated below.

Part 1: Highlight Matching Phrases

Problem Statement

Implement a function:

def highlight_matches(document: str, sources: list[str]) -> str:
    pass

A source matches when its complete text occurs at word boundaries in document.

For the reference implementation in this write-up, assume:

• Matching is case-sensitive. Whether matching should ignore case is left open, so confirm with the interviewer.
• Every source is non-empty and begins and ends with a letter or digit.
• A source must match at word boundaries. For example, "blue" does not match the "blue" inside "blueprint".
• The same source may occur multiple times.
• Matches may overlap.
• Some interviewers also ask to merge adjacent matches. This implementation treats matches separated only by whitespace as adjacent.
• The original document, including its capitalization, punctuation, and whitespace, must be preserved.

Return the document with every merged matched region wrapped in <yellow> and </yellow> tags.

Example

document = "The quick brown fox jumps over the lazy dog."
sources = ["quick brown", "brown fox jumps"]

highlight_matches(document, sources)
# Returns:
# "The <yellow>quick brown fox jumps</yellow> over the lazy dog."

The two matches overlap on "brown", so they become one highlighted interval spanning "quick brown fox jumps".

Edge Cases

• document = "blueprint", sources = ["blue"] produces no highlight because "blue" is not a complete word.
• A source that appears multiple times creates multiple matches.
• Duplicate source strings are still separate sources because citations use their positions in the input list.
• If no source matches, return the original document unchanged.

Part 2: Add Global Counts And Citations

Follow-Up Prompt

Extend the solution so that each highlighted region identifies the sources that contributed to it.

For every source:

1. Count all of its valid occurrences across the entire document. Counts are global, not limited to a single highlighted region.
2. For each merged region, collect the indices of the sources with at least one match inside that region.
3. Order a region's source indices by descending global occurrence count. Tie-breaking is left open; this write-up breaks ties by smaller source index.
4. Append each zero-based source index after the closing tag using the format [source_id].

Return the tagged document, the global count for every source, and the citations attached to each highlighted region.

from dataclasses import dataclass

@dataclass
class CitationResult:
    tagged_document: str
    counts: dict[int, int]
    citations: list[list[int]]

def highlight_with_citations(
    document: str,
    sources: list[str],
) -> CitationResult:
    pass

Example

document = "The quick brown fox jumps over the quick blue fox."
sources = [
    "quick brown",      # source 0, one occurrence
    "brown fox jumps",  # source 1, one occurrence
    "quick",            # source 2, two occurrences
    "blue",             # source 3, one occurrence
]

result = highlight_with_citations(document, sources)

result.tagged_document
# "The <yellow>quick brown fox jumps</yellow>[2][0][1] over the "
# "<yellow>quick blue</yellow>[2][3] fox."

result.counts
# {0: 1, 1: 1, 2: 2, 3: 1}

result.citations
# [[2, 0, 1], [2, 3]]

Source 2 is listed first in both citation groups because it occurs twice globally. Sources with equal counts are ordered by their input index.