Given two sorted arrays A and B, of sizes m and n respectively. Define s = a + b, where a is one element from A and b is one element from B. Find the Kth smallest s out of all possible s'.

Assumptions

• A is not null and A is not of zero length, so as B
• K > 0 and K < = m * n

Examples

A = {1, 3, 5}, B = {4, 8}

• 1st smallest s is 1 + 4 = 5
• 2nd smallest s is 3 + 4 = 7
• 3rd, 4th smallest s are 9 (1 + 8, 4 + 5)
• 5th smallest s is 3 + 8 = 11

Solution: 记得去重

  public int kthSum(int[] A, int[] B, int k) {
    boolean[][] visited = new boolean[A.length][B.length];
    PriorityQueue<Node> pq = new PriorityQueue<>(1, new MyComparator());
    pq.offer(new Node(0, 0, A[0] + B[0]));
    visited[0][0] = true;
    for (int i = 1; i < k; i++) {
      Node cur = pq.poll();
      if (cur.i + 1 < A.length && !visited[cur.i + 1][cur.j]) {
        pq.offer(new Node(cur.i + 1, cur.j, A[cur.i + 1] + B[cur.j]));
        visited[cur.i + 1][cur.j] = true;
      }
      if (cur.j + 1 < B.length && !visited[cur.i][cur.j + 1]) {
        pq.offer(new Node(cur.i, cur.j + 1, A[cur.i] + B[cur.j + 1]));
        visited[cur.i][cur.j + 1] = true;
      }
    }
    return pq.peek().sum;
  }
  class Node {
    int i, j, sum;
    public Node(int i1, int i2, int i3) {
      i = i1;
      j = i2;
      sum = i3;
    }
  }
  class MyComparator implements Comparator<Node> {
    public int compare(Node a, Node b) {
        return a.sum - b.sum;
    }
  }