Binary Search, Priority Queue, Heap
Find the kth smallest number in a row and column sorted matrix.
Example
Given k = 4 and a matrix:
[
[1 ,5 ,7],
[3 ,7 ,8],
[4 ,8 ,9],
]
return 5
Challenge
O(k log n), n is the maximal number in width and height.
Tags
Heap Priority Queue Matrix
Related Problems
Hard Kth Smallest Sum In Two Sorted Arrays
Medium Kth Largest Element
Analysis
Heap
寻找第k小的数，可以联想到转化为数组后排序，不过这样的时间复杂度较高：O(n^2 log n^2) + O(k).
进一步，换种思路，考虑到堆（Heap）的特性，可以建立一个Min Heap，然后poll k次，得到第k个最小数字。不过这样的复杂度仍然较高。
考虑到问题中矩阵本身的特点：排过序，那么可以进一步优化算法。
[1 ,5 ,7],
[3 ,7 ,8],
[4 ,8 ,9],
因为行row和列column都已排序，那么matrix中最小的数字无疑是左上角的那一个，坐标表示也就是(0, 0)。寻找第2小的数字，也就需要在(0, 1), (1, 0)中得出；以此类推第3小的数字，也就要在(0, 1), (1, 0), (2, 0), (1, 1), (0, 2)中寻找。
在一个数字集合中寻找最大（Max）或者最小值（Min），很快可以联想到用Heap，在Java中的实现是Priority Queue，它的pop，push操作均为O(logn)，而top操作，得到堆顶仅需O(1)。
从左上(0, 0)位置开始往右/下方遍历，使用一个Hashmap记录visit过的坐标，把候选的数字以及其坐标放入一个大小为k的heap中（只把未曾visit过的坐标放入heap），并且每次放入前弹出掉（poll）堆顶元素，这样最多会添加（push）2k个元素。时间复杂度是O(klog2k)，也就是说在矩阵自身特征的条件上优化，可以达到常数时间的复杂度，空间复杂度也为O(k)，即存储k个候选数字的Priority Queue （Heap）。
Binary Search
Ref: https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/85173/Share-my-thoughts-and-Clean-Java-Code
Solution
Heap (Priority Queue) with Visited Matrix (17 ms 53.04% AC)
class Number {
public int x, y, val;
public Number(int x, int y, int val) {
this.x = x;
this.y = y;
this.val = val;
}
}
class NumberComparator implements Comparator<Number> {
public int compare(Number a, Number b) {
return a.val - b.val;
}
}
public class Solution {
private boolean isValid(int x, int y, int[][] matrix, boolean[][] visited) {
if (x < matrix.length && y < matrix[x].length && !visited[x][y]) {
return true;
}
return false;
}
int[] dx = new int[] {0, 1};
int[] dy = new int[] {1, 0};
/**
* @param matrix: a matrix of integers
* @param k: an integer
* @return: the kth smallest number in the matrix
*/
public int kthSmallest(int[][] matrix, int k) {
// Validate input
if (matrix == null || matrix.length == 0) {
return -1;
}
if (matrix.length * matrix[0].length < k) {
return -1;
}
// Define min heap
PriorityQueue<Number> heap = new PriorityQueue<Number>(k, new NumberComparator());
heap.add(new Number(0, 0, matrix[0][0]));
// Define visited matrix
boolean[][] visited = new boolean[matrix.length][matrix[0].length];
visited[0][0] = true;
for (int i = 0; i < k - 1; i++) {
Number smallest = heap.poll();
for (int j = 0; j < 2; j++) {
// Next coordinates
int nx = smallest.x + dx[j];
int ny = smallest.y + dy[j];
if (isValid(nx, ny, matrix, visited)) {
visited[nx][ny] = true;
heap.add(new Number(nx, ny, matrix[nx][ny]));
}
}
}
return heap.peek().val;
}
}
Min Heap - Add to min heap, then poll k - 1 times - (56 ms 21.96 % AC)
class Solution {
class Pos {
int i, j, val;
public Pos (int i, int j, int val) {
this.i = i;
this.j = j;
this.val = val;
}
}
public int kthSmallest(int[][] matrix, int k) {
int m = matrix.length;
int n = matrix[0].length;
Queue<Pos> minHeap = new PriorityQueue<>((o1, o2) -> o1.val - o2.val);
for (int i = 0; i < m; i++) {
minHeap.offer(new Pos(i, 0, matrix[i][0]));
}
for (int i = 0; i < k - 1; i++) {
Pos top = minHeap.poll();
if (top.j + 1 < n) {
minHeap.offer(new Pos(top.i, top.j + 1, matrix[top.i][top.j + 1]));
}
}
Pos kth = minHeap.peek();
return kth.val;
}
}
Binary Search
public class Solution {
public int kthSmallest(int[][] matrix, int k) {
int lo = matrix[0][0], hi = matrix[matrix.length - 1][matrix[0].length - 1] + 1;//[lo, hi)
while(lo < hi) {
int mid = lo + (hi - lo) / 2;
int count = 0, j = matrix[0].length - 1;
for(int i = 0; i < matrix.length; i++) {
while(j >= 0 && matrix[i][j] > mid) j--;
count += (j + 1);
}
if(count < k) lo = mid + 1;
else hi = mid;
}
return lo;
}
}
Reference