Heap Frontier Pair Search

Description

Two sorted arrays of n values give n^2 pair sums; you want only the k largest. A matrix whose rows and columns are sorted has n^2 values; you want the k-th smallest. Generating everything costs O(n^2) or more - but the sorted structure means the next answer is always adjacent to answers already found.

This skill covers best-first search over a sorted grid: start at the best corner, keep a heap of frontier cells, and after each pop add only the neighbors that could be next. It also covers when a visited set is needed and when seeding by row avoids it, and the binary-search-on-value alternative for large k.

We'll go from simple to more technical:

1. The pattern: the k largest pair sums, starting at the corner and pushing two neighbors per pop.
2. A second flavor: the k-th smallest in a row- and column-sorted matrix, seeded with each row's head.
3. Where top-k combinations show up (ranking, recommendations, top-k joins).
4. What breaks it: no visited set, the wrong corner, treating the matrix as sorted, the sort's cost.

The approach

int[] maxCombinations(int[] a, int[] b, int k) {
    int[] x = a.clone(), y = b.clone();
    Arrays.sort(x);
    Arrays.sort(y);
    int n = x.length, m = y.length;
    PriorityQueue<int[]> heap = new PriorityQueue<>((p, q) -> Integer.compare(q[0], p[0]));   // max-heap {sum, i, j}
    Set<Long> seen = new HashSet<>();
    heap.offer(new int[]{x[n - 1] + y[m - 1], n - 1, m - 1});      // the largest pair is the last of each
    seen.add((long) (n - 1) * m + (m - 1));
    int[] out = new int[k];
    for (int t = 0; t < k; t++) {
        int[] top = heap.poll();
        out[t] = top[0];
        int i = top[1], j = top[2];
        if (i > 0 && seen.add((long) (i - 1) * m + j)) heap.offer(new int[]{x[i - 1] + y[j], i - 1, j});
        if (j > 0 && seen.add((long) i * m + j - 1)) heap.offer(new int[]{x[i] + y[j - 1], i, j - 1});
    }
    return out;
}

import heapq


def max_combinations(a: list[int], b: list[int], k: int) -> list[int]:
    x, y = sorted(a), sorted(b)
    n, m = len(x), len(y)
    heap = [(-(x[n - 1] + y[m - 1]), n - 1, m - 1)]   # max-heap via negated sums
    seen = {(n - 1, m - 1)}                             # a set of (i, j) tuples - no index encoding needed
    out = []
    for _ in range(k):
        neg, i, j = heapq.heappop(heap)
        out.append(-neg)
        if i > 0 and (i - 1, j) not in seen:
            seen.add((i - 1, j))
            heapq.heappush(heap, (-(x[i - 1] + y[j]), i - 1, j))
        if j > 0 and (i, j - 1) not in seen:
            seen.add((i, j - 1))
            heapq.heappush(heap, (-(x[i] + y[j - 1]), i, j - 1))
    return out

The pairs form a grid where sums shrink as either index decreases. The best pair is the corner (n - 1, m - 1), and every other pair is reachable by stepping down one index at a time - so after each pop, its two "next-best" neighbors join the frontier. A pair can be reached from two neighbors ((i, j) from (i + 1, j) and from (i, j + 1)); without the visited set it's pushed twice and can appear twice in the answer - on the hidden test, 63 shows up twice and the real fifth sum, 56, is lost. Cost: O(n log n) to sort, then O(k log k) for the heap.

A second flavor: k-th smallest in a sorted matrix

A matrix with sorted rows and sorted columns isn't sorted in reading order ([[1, 5], [2, 6]] reads 1, 5, 2, 6), but the same frontier works from the smallest corner. Seed the heap with the head of each row; after popping a cell, push its right neighbor. Each row is entered once, so no visited set is needed:

int kthSmallest(int[][] matrix, int k) {
    int n = matrix.length;
    PriorityQueue<int[]> heap = new PriorityQueue<>((p, q) -> Integer.compare(matrix[p[0]][p[1]], matrix[q[0]][q[1]]));
    for (int r = 0; r < Math.min(n, k); r++) heap.offer(new int[]{r, 0});
    for (int i = 1; i < k; i++) {
        int[] top = heap.poll();
        if (top[1] + 1 < n) heap.offer(new int[]{top[0], top[1] + 1});
    }
    int[] kth = heap.peek();
    return matrix[kth[0]][kth[1]];
}

import heapq


def kth_smallest(matrix: list[list[int]], k: int) -> int:
    n = len(matrix)
    heap = [(matrix[r][0], r, 0) for r in range(min(n, k))]   # the head of each row
    heapq.heapify(heap)
    for _ in range(k - 1):
        _, r, c = heapq.heappop(heap)
        if c + 1 < n:
            heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
    return heap[0][0]

That's O(k log n). The alternative is binary search on the answer: count values <= mid with a staircase walk from the top-right corner in O(n), and find the smallest value whose count reaches k - O(n log(max - min)), independent of k.

Where you'll see it

Recommendation and search systems combine two ranked lists (candidates x jobs, products x offers) and need only the top few combined scores; database engines compute top-k joins the same way without materializing the full cross product. The same best-first frontier is Dijkstra's algorithm in miniature.

Which tool for which job?

Common mistakes

Check yourself

Try to answer before opening each one - recalling beats re-reading.

1. Why does the largest-pair-sum search need a visited set?

2. Why is the method O(n log n + k log k) and not just O(k log k)?

3. Why doesn't the sorted-matrix version need a visited set?

Key takeaways

• Sorted along both axes: start at the best corner and expand the frontier with a heap.
• Mark cells that have two parents as visited, or seed per row so each cell has one parent.
• Largest pair sums: O(n log n) sort + O(k log k) heap.
• Row- and column-sorted isn't row-major sorted; heap O(k log n) or binary search on value.