[LintCode] 465 Kth Smallest Sum in Two Sorted Arrays 解题报告

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


思路
我们把这两个数组一个当横坐标一个当纵坐标看。
       1    7    11
2     3    9    13
4     5    11  15
6     7    13  17
看出名堂了没有。。。就是在这坨红色的数字里找第K小。。。
这不就是Kth Smallest Number in Sorted Matrix这题嘛。。。


Code

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 {
        int row;
        int col;
        int val;
        public Node(int row, int col, int val) {
            this.row = row;
            this.col = col;
            this.val = val;
        }
    } 
     
    public int kthSmallestSum(int[] A, int[] B, int k) {
        // Write your code here
        
        Queue<Node> heap = new PriorityQueue<Node>(k, new Comparator<Node>() {
           public int compare(Node a, Node b) {
               return a.val - b.val;
           } 
        });
        
        for (int row = 0; row < B.length; row++) {
            int col = 0;
            heap.add(new Node(row, 0, A[0] + B[row]));
        }
        
        for (int i = 0; i < k; i++) {
            Node tmp = heap.poll();
            int row;
            int col;
            int val;
            row = tmp.row;
            col = tmp.col;
            val = tmp.val;
            
            if (i == k - 1) {
                return val;
            }
            
            if (col + 1 < A.length) {
                heap.add(new Node(row, col + 1, A[col + 1] + B[row]));
            }
        }
        
        return -1;
    }
}