Find the top k sums of two sorted arrays

Question

You are given two sorted arrays, of sizes n and m respectively. Your task (should you choose to accept it), is to output the largest k sums of the form a[i]+b[j].

A O(k log k) solution can be found here. There are rumors of a O(k) or O(n) solution. Does one exist?

Answer 1

I found the responses at your link mostly vague and poorly structured. Here's a start with a O(k * log(min(m, n))) O(k * log(m + n)) O(k * log(k)) algorithm.

Suppose they are sorted decreasing. Imagine you computed the m*n matrix of the sums as follows:

for i from 0 to m
    for j from 0 to n
        sums[i][j] = a[i] + b[j]

In this matrix, values monotonically decrease down and to the right. With that in mind, here is an algorithm which performs a graph search through this matrix in order of decreasing sums.

q : priority queue (decreasing) := empty priority queue
add (0, 0) to q with priority a[0] + b[0]
while k > 0:
    k--
    x := pop q
    output x
    (i, j) : tuple of int,int := position of x
    if i < m:
        add (i + 1, j) to q with priority a[i + 1] + b[j]
    if j < n:
        add (i, j + 1) to q with priority a[i] + b[j + 1]

Analysis:

1. The loop is executed k times.
   1. There is one pop operation per iteration.
   2. There are up to two insert operations per iteration.
2. The maximum size of the priority queue is O(min(m, n)) O(m + n) O(k).
3. The priority queue can be implemented with a binary heap giving log(size) pop and insert.
4. Therefore this algorithm is O(k * log(min(m, n))) O(k * log(m + n)) O(k * log(k)).

Note that the general priority queue abstract data type needs to be modified to ignore duplicate entries. Alternately, you could maintain a separate set structure that first checks for membership in the set before adding to the queue, and removes from the set after popping from the queue. Neither of these ideas would worsen the time or space complexity.

I could write this up in Java if there's any interest.

Edit: fixed complexity. There is an algorithm which has the complexity I described, but it is slightly different from this one. You would have to take care to avoid adding certain nodes. My simple solution adds many nodes to the queue prematurely.

Answer 2

private static class FrontierElem implements Comparable<FrontierElem> {
    int value;
    int aIdx;
    int bIdx;

    public FrontierElem(int value, int aIdx, int bIdx) {
        this.value = value;
        this.aIdx = aIdx;
        this.bIdx = bIdx;
    }

    @Override
    public int compareTo(FrontierElem o) {
        return o.value - value;
    }

}

public static void findMaxSum( int [] a, int [] b, int k ) {
    Integer [] frontierA = new Integer[ a.length ];
    Integer [] frontierB = new Integer[ b.length ];
    PriorityQueue<FrontierElem> q = new PriorityQueue<MaxSum.FrontierElem>();
    frontierA[0] = frontierB[0]=0;
    q.add( new FrontierElem( a[0]+b[0], 0, 0));
    while( k > 0 ) {
        FrontierElem f = q.poll();
        System.out.println( f.value+"    "+q.size() );
        k--;
        frontierA[ f.aIdx ] = frontierB[ f.bIdx ] = null;
        int fRight = f.aIdx+1;
        int fDown = f.bIdx+1;
        if( fRight < a.length && frontierA[ fRight ] == null ) {
            q.add( new FrontierElem( a[fRight]+b[f.bIdx], fRight, f.bIdx));
            frontierA[ fRight ] = f.bIdx;
            frontierB[ f.bIdx ] = fRight;
        }
        if( fDown < b.length && frontierB[ fDown ] == null ) {
            q.add( new FrontierElem( a[f.aIdx]+b[fDown], f.aIdx, fDown));
            frontierA[ f.aIdx ] = fDown;
            frontierB[ fDown ] = f.aIdx;
        }
    }
}

The idea is similar to the other solution, but with the observation that as you add to your result set from the matrix, at every step the next element in our set can only come from where the current set is concave. I called these elements frontier elements and I keep track of their position in two arrays and their values in a priority queue. This helps keep the queue size down, but by how much I've yet to figure out. It seems to be about sqrt( k ) but I'm not entirely sure about that.

(Of course the frontierA/B arrays could be simple boolean arrays, but this way they fully define my result set, This isn't used anywhere in this example but might be useful otherwise.)