Kth Smallest Element in a Sorted Matrix

Given anxnmatrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note:
You may assume k is always valid, 1 ≤ k ≤ n^2.

Solution #1:

Maintain a max heap of size k. All all numbers into heap and return the top one.

public int kthSmallest(int[][] matrix, int k) {
    PriorityQueue<Integer> pq = new PriorityQueue<>(new Comparator<Integer>() {
        @Override
        public int compare(Integer a, Integer b) {
            return b - a;
        }
    });
    for (int i = 0; i < matrix.length; i++) {
        for (int j = 0; j < matrix[0].length; j++) {
            pq.offer(matrix[i][j]);
            if (pq.size() > k) {
                pq.poll();
            }
        }
    }
    return pq.poll();
}

Solution #2:

The top left corner is the lower bound of the entire matrix, and the bottom right corner is the upper bound of the entire matrix. Guess by using binary search between min and max. Count how many numbers smaller or equal to mid. Same approach in Search a 2D Matrix II.

public int kthSmallest(int[][] matrix, int k) {
    int left = matrix[0][0],
        right = matrix[matrix.length - 1][matrix[0].length - 1];
    while (left < right) {
        int mid = left + (right - left) / 2;
        int count = binarySearch(matrix, mid);
        if (count < k) {
            left = mid + 1;
        } else {
            right = mid;
        }
    }
    return right;
}

private int binarySearch(int[][] matrix, int target) {
    int row   = matrix.length - 1,
        col   = 0,
        count = 0;
    while (row >= 0 && col < matrix[0].length) {
        if (matrix[row][col] <= target) {
            count = count + row + 1;
            col++;
        } else {
            row--;
        }
    }
    return count;
}