Kth Smallest Number in Sorted Matrix

`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

``````[1 ,5 ,7],
[3 ,7 ,8],
[4 ,8 ,9],``````

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());

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

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

Last updated