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:

  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.

假设 (i, j) 指 A[i] + A[j], 那么(i, j)的下一个元素为(i + 1, j)或(i, j + 1).

  1. 将(0, 0)加入最小堆，并用HashSet记录位置，
  2. 将堆顶元素出堆，得到其位置(i，j), 并将(i + 1, j)和(i, j + 1)入堆，HashSet记录位置. 入堆的时候要检查该元素是否入过堆，及坐标不得越界
  3. 重复step 2, k - 1次
  4. 返回堆顶
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Node implements Comparable<Node> {
    int i;
    int j;
    int sum;
    public Node(int i, int j, int sum) {
        this.i = i;
        this.j = j;
        this.sum = sum;
    }
    @Override
    public int compareTo(Node that) {
        return this.sum - that.sum;
    }
}
public int kthSmallestSum(int[] A, int[] B, int k) {
    if (A.length == 0 && B.length == 0)
        return 0;
    if (A.length == 0)
        return B[k - 1];
    if (B.length == 0)
        return A[k - 1];
    PriorityQueue<Node> pq = new PriorityQueue<>();
    HashSet<String> set = new HashSet<>();
    pq.add(new Node(0, 0, A[0] + B[0]));
    set.add(0 + "-" + 0);
    for (int i = 0; i < k - 1; i++) {
        Node n = pq.poll();
        if (!set.contains((n.i + 1) + "-" + n.j) && n.i + 1 < A.length && n.j < B.length) {
            pq.offer(new Node(n.i + 1, n.j, A[n.i + 1] + B[n.j]));
            set.add((n.i + 1) + "-" + n.j);
        }
        if (!set.contains(n.i + "-" + (n.j + 1)) && n.i < A.length && n.j + 1 < B.length) {
            pq.offer(new Node(n.i, n.j + 1, A[n.i] + B[n.j + 1]));
            set.add(n.i + "-" + (n.j + 1));
        }
    }
    return pq.peek().sum;
}