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):

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;
    }
}