40. Two-Sum on a Token Budget

Problem

A batch-inference scheduler wants to pack exactly two queued prompts into one forward pass so their combined token count exactly fills the available context-window budget, with no waste and no overflow.

Given the queue's token counts and the budget, find the two prompts to pair. Return the pair of indices (i, j) with i < j such that token_counts[i] + token_counts[j] == budget, or None if no such pair exists. If multiple pairs qualify, return the one found first scanning left to right by j. This must run in O(n) time using a hash map, not the O(n²) all-pairs check.

Source: src/40_two_sum_token_budget.py

def two_prompts_within_budget(token_counts: list[int], budget: int) -> tuple[int, int] | None:
    ...

Examples:
>>> two_prompts_within_budget([120, 45, 900, 260, 40], 300)
(3, 4)

>>> two_prompts_within_budget([100, 200], 500)
None

Step-by-Step Approach

1. Recognize that for each prompt, the only thing that matters is whether its "complement" (budget - token_counts[i]) has already appeared earlier in the queue.
2. Maintain a hash map from token count to the earliest index at which that count was seen, built up as you scan left to right.
3. At each index j, before recording anything, compute the complement needed to hit the budget and check whether it's already a key in the map.
4. If the complement is present, you've found your pair immediately: the stored index and the current index j, with the stored one guaranteed smaller since it was recorded on an earlier iteration.
5. If not, record the current count in the map (only if that value hasn't been seen yet, so ties resolve to the earliest index) and move on.
6. If the scan finishes with no match found, no valid pair exists in the queue, so return None.

The key insight is that checking "has the complement been seen" is an O(1) hash lookup instead of an O(n) inner scan, which is what collapses the classic O(n²) all-pairs approach down to a single O(n) pass.

Reference solution

def two_prompts_within_budget(token_counts: list[int], budget: int) -> tuple[int, int] | None:
    # single pass, hash map from token count -> earliest index seen, O(n) time/space
    seen: dict[int, int] = {}
    for j, count in enumerate(token_counts):
        complement = budget - count
        if complement in seen:
            return (seen[complement], j)
        # only record the first occurrence of each value so ties resolve to earliest indices
        if count not in seen:
            seen[count] = j
    return None

Key Functions & Tricks

• enumerate(token_counts) — single left-to-right pass, tracking both value and index
• Complement lookup via hash map — turns "does some earlier element sum with this one to the target" into an O(1) membership check instead of a nested loop
• complement = budget - count — the value that must already exist in seen for a match
• if count not in seen guard — keeps the earliest index for each distinct value, so returned pairs are deterministic on duplicates
• Early return on first match — avoids scanning the rest of the queue once a valid pair is found

How to Recognize This Pattern