Kth Smallest Sum In Two Sorted Arrays

Hard
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.

我自己用的maxHeap的暴力解：

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
     */
    public int kthSmallestSum(int[] A, int[] B, int k) {
        // write your code here
        int n = A.length;
        int m = B.length;
        PriorityQueue pq = new PriorityQueue<>(n + m, maxHeap);
        for (int i = 0; i < n; i++){
            for (int j = 0; j < m; j++){
                int sum = A[i] + B[j];
                pq.add(sum);
            }
        }
        while (pq.size() > k){
            pq.poll();
        }
        return pq.poll(); 
    }
    private Comparator maxHeap = new Comparator(){
        public int compare(Integer a, Integer b){
            return b - a;
        }
    };
}

但是这个方法time complexity是O(N2)的，因为用了两层for循环把所有的元素都加到了heap里。而Challenge里要求的是
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.

把问题转化成很类似kth smallest number in sorted matrix. 我们从A[0],B[0]开始放进heap, 很明显每次放进去都是最小的。然后每次poll()出来的时候，加上它的“坐标”的(x + 1, y)和（x , y + 1)，这样poll（）掉k - 1个，拿出来的就是最小的k - 1个sum. 最后再peek()的就是Kth Smallest Sum.

Submit 1: (一直没发现错在了哪儿）

Screen Shot 2017-09-29 at 2.55.41 PM.png

后来请教群里的人发现是String的处理细节出了问题：1 + "," + 3 + 1 得到的是 1,31而不是1，4

AC:

Screen Shot 2017-09-29 at 3.12.30 PM.png

public class Solution {
    public class Point{
        int x;
        int y;
        int sum;
        public Point(int x, int y, int sum){
            this.x = x;
            this.y = y;
            this.sum = sum;
        }
    }
    /*
     * @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
     */
    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 < 0){
            return - 1;
        }
        int n = A.length;
        int m = B.length;
        PriorityQueue pq = new PriorityQueue<>(k, minHeap);
        HashSet visited = new HashSet<>();
        Point p = new Point(0, 0, A[0] + B[0]);
        pq.offer(p);
        visited.add(p.x + "," + p.y);
        for (int i = 0; i < k - 1; i++){
            Point pt = pq.poll();
            if (pt.x + 1 < n && pt.y < m && !visited.contains((pt.x+1) + "," + pt.y) ){
                visited.add((pt.x+1) + "," + pt.y);
                pq.offer(new Point(pt.x + 1, pt.y, A[pt.x + 1] + B[pt.y]));
            }
            if (pt.x < n && pt.y + 1 < m && !visited.contains(pt.x + "," + (pt.y+1))){
                visited.add(pt.x + "," + (pt.y+1));
                pq.offer(new Point(pt.x, pt.y + 1, A[pt.x] + B[pt.y + 1]));
            }
        }
        return pq.peek().sum;
    }
    
    private Comparator minHeap = new Comparator(){
        public int compare(Point a, Point b){
            return a.sum - b.sum;
        }
    };
}

注意一下时间复杂度的分析,显示PriorityQueue里面各个操作的Big O:

What is time complexity for offer, poll and peek methods in priority queue Java?
Answer: Time complexity for the methods offer & poll is O(log(n)) and for the peek() it is Constant time O(1).

NOTES:

 In Java programming, Java Priority Queue is implemented using Heap         
 Data Structures and Heap has O(log(n)) time complexity to insert and 
 delete element.
 Offer() and add() methods are used to insert the element in the queue.
 Poll() and remove() is used to delete the element from the queue.
 Element retrieval methods i.e. peek() and element(), that are used to 
 retrieve elements from the head of the queue is constant time i.e. O(1).
 contains(Object)method that is used to check if a particular element is 
 present in the queue, have leaner time complexity i.e. O(n).

所以每一次poll()需要logn，这里的n指的是pq里的元素个数，我们直到所有的sum个数只能是min(m,n) 同时pq的size limit是k, 所以我们每一次poll()需要log min(m,n,k), for循环k次，所以时间复杂度是O(k.log(min(m,n,k))