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 {
        public int x;
        public int y;
        public int val;
        public Node(int x, int y, int val) {
            this.x = x;
            this.y = y;
            this.val = val;
        }
    }

    public int kthSmallestSum(int[] A, int[] B, int k) {
        int m = A.length;
        int n = B.length;
        PriorityQueue<Node> pq = new PriorityQueue<Node>(m * n, (a, b) -> (a.val - b.val));
        boolean[][] visited = new boolean[m][n];
        // matrix[i][j] = A[i] + B[j];
        pq.offer(new Node(0, 0, A[0] + B[0]));

        int count = 0;
        int[][] dirs = {{0,1},{1,0}};
        while(count < k) {
            Node node = pq.poll();
            if(visited[node.x][node.y]) {
                continue;
            }
            visited[node.x][node.y] = true;
            count++;
            if(count == k) {
                return node.val;
            }
            for(int[] dir: dirs) {
                int nx = node.x + dir[0];
                int ny = node.y + dir[1];
                if(0 <= nx && nx < m && 0 <= ny && ny < n && !visited[nx][ny]) {
                    pq.offer(new Node(nx, ny, A[nx] + B[ny]));
                }
            }
        }
        return -1;
    }
}