Problem link: https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/
constraint:
n == matrix.length
n == matrix[i].length ==> **count of rows == count of columns**
1 <= n <= 300
-109 <= matrix[i][j] <= 109
All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
1 <= k <= n2 ==> k is always valid
Idea
- Given that the question is asking for Kth smallest number, we should think about using Priority Heap. Since all the numbers within a row/col are non-decreasing, we could think of all the columns as separate list of sorted numbers, and we have to find the Kth smallest one among them. we could start by adding into the heap all the numbers at the head of each lists, (which also means all numbers within the first row). Everytime we pop the top element from the heap (also the smallest so far) and then push its next number within that list. Repeat this for K times, then we will get the results we want. The element in the heap should contains info about which row/col the number is in.
class Solution {
public int kthSmallest(int[][] matrix, int k) {
PriorityQueue<Node> pq = new PriorityQueue<>();
for (int c = 0; c < matrix[0].length; c++) {
pq.offer(new Node(0, c, matrix[0][c]));
}
int res = Integer.MAX_VALUE;
for (int i = 0; i < k; i++) {
Node minNode = pq.poll();
res = minNode.val;
if (minNode.row + 1 < matrix.length) {
Node nextNode = new Node(minNode.row + 1, minNode.col, matrix[minNode.row + 1][minNode.col]);
pq.offer(nextNode);
}
}
return res;
}
private class Node implements Comparable<Node> {
int row;
int col;
int val;
Node(int row, int col, int val) {
this.row = row;
this.col = col;
this.val = val;
}
@Override
public int compareTo(Node other) {
return this.val - other.val;
}
}
}
- Time: Initialzing the heap takes O(nlogn). Finding the smallest element for K times takes O(klogn). But since 1 <= k <= n^2, the overall time should also be O(nlogn).
- Space: O(n) as the heap need to store n elements.
- Binary search