Maximum Sum Combination
Maximum Sum Combination: from arrays A and B, return the k largest sums A[i] + B[j] over all pairs.
- 1 <= n <= 10⁵
- 1 <= k <= n
- 1 <= A[i], B[i] <= 10⁵
Intuition
Maximum sum combination takes two arrays and asks for the k largest values of A[i] + B[j] across all pairs. There are n² pairs, and computing them all then sorting is O(n² log n) — far too much when k is small.
Sort both arrays in descending order and picture the sums as a grid, where cell (i, j) holds A[i] + B[j]. Because both axes are sorted, values decrease as you move right and as you move down. The largest sum is at the top-left corner, and the k largest form a staircase-shaped region hugging that corner.
So explore best-first instead of exhaustively. Start at (0, 0), the guaranteed maximum. Once a cell is taken, the next candidates are its two neighbours — (i+1, j) and (i, j+1) — because every other cell is dominated by one of them. A max-heap holding the frontier always yields the next largest sum:
- Pop the best cell, emit it, push its two neighbours, repeat k times.
One subtlety makes or breaks it. Cell (1, 1) can be reached from (0, 1) and from (1, 0), so without protection it enters the heap twice and gets emitted twice. A visited set keyed by (i, j) is required, not an optimisation. With it, each emit pushes at most two cells, so the heap stays around O(k) and only about k cells are ever touched.
Both arrays sorted descending means the best pair is (0, 0), and the next best is always a neighbour of something already taken. That's the k-way frontier pattern: keep a heap of candidates, pop the best, and push only the cells adjacent to it. The seen set is essential because two different pops can reach the same cell.
Approach
Before reading on: price up what the direct approach costs here, then ask why only the extreme element matters and the rest need no ordering. Aim for O(n log n + k log k) time and O(k) space.
Sort both arrays descending
Sorting makes the sum grid monotone — values decrease rightward and downward. That monotonicity is what guarantees the frontier contains the next best sum, and without it the heap approach has no basis.
Seed the heap with the corner
Push (A[0] + B[0], 0, 0), which is provably the largest sum since both components are maximal. Mark (0, 0) visited at the same moment, so the marking and the pushing never drift apart.
Pop the best and emit it
Each pop yields the largest sum in the frontier, which is the next largest overall — every unexplored cell is dominated by something already in the heap. Repeat k times to collect the answer in descending order.
Push the two neighbours
After popping (i, j), push (i+1, j) and (i, j+1) if they are in bounds and unvisited. These are the only cells that can become the new maximum, since anything further away is dominated by one of them.
Guard with a visited set
Without it (i+1, j+1) is reachable by two paths and gets emitted twice, producing duplicate sums and a wrong answer. Mark a cell at push time, not at pop time, or the duplicate is already inside the heap before you check.
Cost of the best-first search
Each of the k pops pushes at most two cells, so the heap stays O(k) and the total is O(k log k) after the O(n log n) sort. Compare with O(n²) to generate every sum — the saving is large whenever k is much smaller than n².
Solution & live demo
Common pitfalls
Building all n×m sums
sums = sorted((x + y for x in a for y in b), reverse=True) return sums[:k]
heap = [(-(a[0] + b[0]), 0, 0)] # expand only the frontier
That's O(nm log nm) work and memory for the k largest values. The heap only ever holds the frontier, so the cost is O(k log k) regardless of how large the arrays are.
Omitting the visited set
for ni, nj in ((i + 1, j), (i, j + 1)):
heapq.heappush(heap, (-(a[ni] + b[nj]), ni, nj))if (ni, nj) not in seen:
seen.add((ni, nj))
heapq.heappush(heap, ...)Cell (1, 1) is reachable from both (0, 1) and (1, 0), so without the guard the same sum is pushed twice and appears twice in the output. The set makes each cell enter the heap exactly once.
Forgetting to negate for a max-heap
heap = [(a[0] + b[0], 0, 0)]
heap = [(-(a[0] + b[0]), 0, 0)]
heapq only pops minimums, so storing raw sums yields the k smallest combinations. Values are negated going in and negated again coming out.
Edge cases
Just A[0]+B[0] after sorting — one pop.
Different (i,j) cells may tie; both count, the set dedupes cells not values.