Kth Smallest Sum In Two Sorted Arrays

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

Example 1

Input:
a = [1, 7, 11]
b = [2, 4, 6]
k = 3
Output: 7
Explanation: The sums are [3, 5, 7, 9, 11, 13, 13, 15, 17] and the 3th is 7.
           

Example 2

Input:
a = [1, 7, 11]
b = [2, 4, 6]
k = 4
Output: 9
Explanation: The sums are [3, 5, 7, 9, 11, 13, 13, 15, 17] and the 4th is 9.
           

Example 3

Input:
a = [1, 7, 11]
b = [2, 4, 6]
k = 8
Output: 15
Explanation: The sums are [3, 5, 7, 9, 11, 13, 13, 15, 17] and the 8th is 15.
           

Challenge

Do it in either of the following time complexity:

1. O(k log min(n, m, k)). where n is the size of A, and m is the size of B.
2. O( (m + n) log maxValue). where maxValue is the max number in A and B.

思路：这题你要知道把A的每一个元素加上B的一整条元素构成一行，这样可以构建一个m*n的矩阵，从左往右，从上往下都是递增的。 

a = [1, 7, 11]

b = [2, 4, 6]

matrix[][] =  [1+2, 1+4, 1+6]

                   [7+2, 7+4, 7+6]

                   [11+2, 11+4, 11+6]

这样问题就转换为 Kth Smallest Element in a Sorted Matrix ,这个矩阵没必要建立起来，因为matrix[i][j] = A[i] + A[j],所以记录坐标就可以了。代码跟 Kth Smallest Element in a Sorted Matrix 一模一样；Time Complexity: KlogK

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
     */
     
    private class Node {
        int x;
        int y;
        int value;
        public Node(int x, int y, int value) {
            this.x = x;
            this.y = y;
            this.value = value;
        }
    }
    
    private class myComparator implements Comparator<Node> {
        @Override
        public int compare(Node a, Node b) {
            return a.value - b.value;
        }
    }
    
    public int kthSmallestSum(int[] A, int[] B, int k) {
        int m = A.length; 
        int n = B.length;
        
        boolean[][] visited = new boolean[m][n];
        PriorityQueue<Node> pq = new PriorityQueue<Node>(k, new myComparator());
        pq.add(new Node(0, 0, A[0] + B[0]));
        
        int[] dx = {0, 1};
        int[] dy = {1, 0};
        
        int count = 0;
        int res = 0;
        while(count < k) {
            Node node = pq.poll();
            count++;
            if(count == k) {
                res = node.value;
            }
            for(int i = 0; i < 2; i++) {
                int nx = node.x + dx[i];
                int ny = node.y + dy[i];
                if(nx < m && ny < n && !visited[nx][ny]) {
                    pq.add(new Node(nx, ny, A[nx] + B[ny]));
                    visited[nx][ny] = true;
                }
            }
        }
        return res;
    }
}