Kth Smallest Sum In Two Sorted Arrays 465

Question

Given two integer arrays sorted in ascending order and an integer k. Define sum = a + b, where a is an element from the first array and b is an element from the second one. Find the kth smallest sum out of all possible sums.

Example

Given [1, 7, 11] and [2, 4, 6].

For k = 3, return 7.

For k = 4, return 9.

For k = 8, return 15.

Challenge

Do it in either of the following time complexity:

O(k log min(n, m, k)). where n is the size of A, and m is the size of B.

O( (m + n) log maxValue). where maxValue is the max number in A and B.

Solution

idea：位置为i（A中），j（B中）元素（表示A[i]+B[j]）的下一个元素为位置i+1,j和i,j+1之中的一个。

1. 将A[0]和B[0]的和加入一个最小堆，并将位置0，0加入一个hashset记录加入过堆的位置

2. 将最堆顶元素出队，得到其位置i，j，并将A[i+1]+B[j]和A[i]+B[j+1]加入堆中，其中i＋1和j+1必须是合法的，即不越界，同时i＋1，j和i，j+1没有加入过堆

3. 重复2，出队k次即可

代码如下：

public class Solution {
    /**
     * @param A an integer arrays sorted in ascending order
     * @param B an integer arrays sorted in ascending order
     * @param k an integer
     * @return an integer
     */
    class Node implements Comparable<Node>{
        int i;
        int j;
        int val;
        public Node(int i, int j, int val){
            this.i = i;
            this.j = j;
            this.val = val;
        }

        @Override
        public int compareTo(Node another){
            if(this.val == another.val){
                return 0;
            }
            return this.val > another.val? 1 : -1;
        }
    }

    int[] di = new int[]{1, 0};
    int[] dj = new int[]{0, 1};
    public int kthSmallestSum(int[] A, int[] B, int k) {
        // Write your code here
        if(A == null || B == null || A.length == 0 && B.length == 0 || k > A.length * B.length){
            return 0;
        }else if(A.length == 0){
            return B[k];
        }else if(B.length == 0){
            return A[k];    
        }

        HashSet<String> set = new HashSet<String>();
        PriorityQueue<Node> queue = new PriorityQueue<Node>();
        Node temp = new Node(0, 0, A[0] + B[0]);
        queue.offer(temp);
        set.add(temp.i + "," + temp.j);
        Node next;
        int i = 0;
        int j = 0;
        for(int n = 0; n < k - 1; n++){
            temp = queue.poll();
            for(int m = 0; m < 2; m++){
                i = temp.i + di[m];
                j = temp.j + dj[m];
                next = new Node(i, j, 0);
                if(i >= 0 && i < A.length && j >= 0 && j < B.length && !set.contains(next.i + "," + next.j)){
                    next.val = A[i] + B[j];
                    queue.offer(next);
                    set.add(next.i + "," + next.j);
                }
            }
        }
        return queue.peek().val;
    }
}