☕ DSA Java Articles

Topics › Heaps › Hard Problems

Maximum Sum Combination

Hard TUF ID 568 A2ZSDE

Java source ↗ Edit this article ↗

Problem

You are given two integer arrays a (length n) and b (length m) and an integer k. A sum combination is a[i] + b[j] for one index from each array. Return the k largest sum combinations in non-increasing order. Different index pairs count as different combinations even when their sums are equal.

Constraints in the usual statement: n = m ≤ 10^5, 1 ≤ k ≤ n. The code here also allows n ≠ m and any k ≤ n · m, and it computes sums in long so that large values cannot overflow.

Example 1

a = [3, 2], b = [1, 4], k = 2
answer = [7, 6]          // 3 + 4 = 7, then 2 + 4 = 6

Example 2

a = [1, 4, 2, 3], b = [2, 5, 1, 6], k = 4
answer = [10, 9, 9, 8]   // 4+6, 4+5, 3+6, then 3+5 or 2+6

The two 9s come from different pairs, (4, 5) and (3, 6), so both are listed.

Approach 1: Brute force, all sums sorted

Intuition. There are exactly n · m sum combinations. Compute them all, sort them, and read the k largest from the end. This is clearly correct and works for small inputs, but n = m = 10^5 gives 10^10 sums, far too many to compute, let alone store.

Algorithm.

  1. Fill an array with a[i] + b[j] for every pair, using long arithmetic.
  2. Sort it ascending.
  3. Copy the last k values, from the end backwards, into the result.
static long[] bruteForce(int[] a, int[] b, int k) {
    long[] all = new long[a.length * b.length];
    int idx = 0;
    for (int x : a) {
        for (int y : b) all[idx++] = (long) x + y;
    }
    Arrays.sort(all);
    long[] res = new long[k];
    for (int t = 0; t < k; t++) res[t] = all[all.length - 1 - t];
    return res;
}

Complexity. Time O(nm log(nm)). Space O(nm).

Approach 2: Better, size-k min-heap

Intuition. We only need the top k, so there is no reason to keep all n · m sums. Stream them through a min-heap that holds at most k values. Its root is the smallest of the current top k. A new sum enters only if it beats that root, and the root is then evicted. This is the standard “k largest” pattern.

Algorithm.

  1. For every pair, compute s = a[i] + b[j].
  2. If the heap has fewer than k values, push s. Otherwise, if s is larger than the root, poll the root and push s.
  3. Poll the heap into the result from the back, because the heap releases the smallest value first.
static long[] better(int[] a, int[] b, int k) {
    PriorityQueue<Long> pq = new PriorityQueue<>();          // the k largest sums seen so far
    for (int x : a) {
        for (int y : b) {
            long s = (long) x + y;
            if (pq.size() < k) {
                pq.add(s);
            } else if (s > pq.peek()) {
                pq.poll();
                pq.add(s);
            }
        }
    }
    long[] res = new long[pq.size()];
    for (int t = res.length - 1; t >= 0; t--) res[t] = pq.poll();   // smallest comes out first
    return res;
}

Complexity. Time O(nm log k). Space O(k).

Memory is fixed, but we still look at every pair. To do better, we must avoid generating sums that obviously cannot be in the top k.

Approach 3: Optimal, sorted arrays and a frontier heap

Intuition. Sort both arrays ascending into x and y. Picture the grid of all sums, S(i, j) = x[i] + y[j]. Because both arrays are sorted, the grid is sorted along every row and every column: moving down or right never decreases the sum. For example 2, with x = [1, 2, 3, 4] and y = [1, 2, 5, 6]:

          y=1  y=2  y=5  y=6
x=1        2    3    6    7
x=2        3    4    7    8
x=3        4    5    8    9
x=4        5    6    9   10

The largest sum is the bottom-right corner, S(n-1, m-1). After taking it, the next largest must be one of its two neighbours, S(n-2, m-1) or S(n-1, m-2). Every other cell has one of those two as a larger-or-equal cell between it and the corner.

The same holds at every step. Any cell (i, j) other than the corner has a neighbour (i+1, j) or (i, j+1) whose sum is at least as large, so that neighbour comes out of the heap no later than (i, j) itself. So it is enough to push only the up and left neighbours (i-1, j) and (i, j-1) of each cell we pop. The heap holds a small frontier of candidates, and its maximum is always the next answer.

One cell can be reached from two directions. (i-1, j-1) is the left neighbour of (i-1, j) and the up neighbour of (i, j-1). A seen set of index pairs stops it from being pushed twice. Without the set, the same pair would be reported twice.

Algorithm.

  1. Sort copies x and y ascending.
  2. Push {x[n-1] + y[m-1], n-1, m-1} into a max-heap ordered by sum, and mark (n-1, m-1) as seen.
  3. Repeat k times: poll the top {sum, i, j} and append sum to the result. If i > 0 and (i-1, j) is unseen, mark it and push it. If j > 0 and (i, j-1) is unseen, mark it and push it.
  4. Return the result.

The pair (i, j) is stored in the set as the single long key i · m + j.

static long[] optimal(int[] a, int[] b, int k) {
    int n = a.length, m = b.length;
    int[] x = a.clone(), y = b.clone();
    Arrays.sort(x);
    Arrays.sort(y);
    PriorityQueue<long[]> pq = new PriorityQueue<>((p, q) -> Long.compare(q[0], p[0]));   // {sum, i, j}, largest first
    Set<Long> seen = new HashSet<>();                       // pair (i, j) encoded as i * m + j
    pq.add(new long[]{(long) x[n - 1] + y[m - 1], n - 1, m - 1});
    seen.add((long) (n - 1) * m + (m - 1));
    long[] res = new long[k];
    for (int t = 0; t < k; t++) {
        long[] top = pq.poll();
        res[t] = top[0];
        int i = (int) top[1], j = (int) top[2];
        if (i > 0 && seen.add((long) (i - 1) * m + j)) pq.add(new long[]{(long) x[i - 1] + y[j], i - 1, j});
        if (j > 0 && seen.add((long) i * m + (j - 1))) pq.add(new long[]{(long) x[i] + y[j - 1], i, j - 1});
    }
    return res;
}

Dry run on example 2 (cells written as (i, j)=sum, using the grid above):

step popped pushed result so far
1 (3,3)=10 (2,3)=9, (3,2)=9 10
2 (2,3)=9 (1,3)=8, (2,2)=8 10, 9
3 (3,2)=9 (3,1)=6; (2,2) already seen 10, 9, 9
4 (1,3)=8 or (2,2)=8   10, 9, 9, 8

When two 9s tie, the heap may return them in either order. The list of sums is the same either way, and that is all the problem asks for.

Complexity. Time O(n log n + m log m + k log k). The sorts dominate when k is small. Each pop pushes at most two cells, so the heap and the set never hold more than 2k + 1 entries. Space O(k) beyond the sorted copies.

A variant without the seen set treats each row as a sorted list: seed the heap with (i, m-1) for every i, and after popping (i, j) push only (i, j-1). That is a k-way merge, costing O(n log n + k log n). It is equally valid, and it is the better choice when k is much larger than n.

Edge cases and pitfalls