Lintcode 465. 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

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.

Code (Java):

?

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 45 46 47 48 49 50

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      */     public int kthSmallestSum(int[] A, int[] B, int k) {         PriorityQueue<Node> pq = new PriorityQueue<>(new MyPQComparator());         boolean[][] visited = new boolean[A.length][B.length];         pq.offer(new Node(0, 0, A[0] + B[0]));         visited[0][0] = true;         int[][] dirs = {{0, 1}, {1, 0}};         int ans = 0;         for (int i = 0; i < k; i++) {             Node curr = pq.poll();             if (i == k - 1) {                 ans = curr.val;             }             for (int j = 0; j < 2; j++) {                 int nextRow = curr.row + dirs[j][0];                 int nextCol = curr.col + dirs[j][1];                 if (nextRow < A.length && nextCol < B.length && !visited[nextRow][nextCol]) {                     pq.offer(new Node(nextRow, nextCol, A[nextRow] + B[nextCol]));                     visited[nextRow][nextCol] = true;                 }             }         }         return ans;     } }   class Node {     public int row;     public int col;     public int val;       public Node (int row, int col, int val) {         this.row = row;         this.col = col;         this.val = val;     } }   class MyPQComparator implements Comparator<Node> {     @Override     public int compare(Node a, Node b) {         return a.val - b.val;     } }