Leetcode — 378. Kth Smallest Element in a Sorted Matrix

Anj
10 min readJan 23, 2021

--

Given an n x n matrix where each of the rows and columns is 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.
You may assume k is always valid, i.e., 1 ≤ k ≤ n².

Example
Input: matrix = [[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]], k = 8
Output: 13.

Thoughts:

If this is just a regular n*n matrix, we will just put all the numbers into a 1D array, sort them, and return the kth element. The time complexity will also be O(n²lgn²) since we sort n² elements.

However, we need to take advantage of BOTH the “sorted rows” and “sorted columns”.

That is to say, matrix[i][j] is definitely smaller than matrix[i+1][j] and matrix[i][j+1].

Example Matrix

This gives us some hint — if we have found the (k-1)the smallest number, then the next smallest number must be appeared in all the previous (k-1) numbers’ right cells or their cells below.

So the question is, how we get the kth smallest number in the matrix?

It is not hard to see the smallest number in the matrix is at (0, 0). Hence, the 2nd smallest number will either be on its right or below, if cell(1, 0) has the 2nd smallest, then we will find the 3rd smallest number among cell(0,1), cell(1,1), and cell(2,0)… etc.

Hence, we kind of have the idea — we will have a min-heap kind of data structure that helps us store the numbers from small to big. We will keep popping up the smallest number among others in the heap. In the meantime when we find the “current” smallest number, we also need to add its right cell and the cell below into the heap since they might be the “next” smallest number. We won’t consider any other numbers that are not in the heap until their neighbor has been chosen.

Since we don’t have a minheap in Java, we will use PriorityQueue to help us do the work. Also, I will create another variable to record the visited positions in case we visit the position more than once (You can use Set, or anything to record the visited positions. I will just use a boolean 2D array since it’s readable). Furthermore, since we not only need the value but also the position of the number, I also create a customizable class. It is fine how you implement it, you can also use an array that contains the value and position as long as the logic is correct.

Below is the 17ms code that beats 38%:

class Solution {
public int kthSmallest(int[][] matrix, int k) {
if(k==1)
return matrix[0][0];
int n = matrix.length;
PriorityQueue<Node> queue = new PriorityQueue<Node>();
boolean[][] visited = new boolean[n][n];
queue.offer(new Node(0,0,matrix[0][0]));
visited[0][0] = true;
while(!queue.isEmpty())
{
Node cur = queue.poll();
k--;
if(k==0)
return cur.val;

int i = cur.i, j = cur.j;
if(i!=n-1 && !visited[i+1][j])
{
queue.offer(new Node(i+1,j,matrix[i+1][j]));
visited[i+1][j] = true;
}
if(j!=n-1 && !visited[i][j+1])
{
queue.offer(new Node(i,j+1,matrix[i][j+1]));
visited[i][j+1] = true;
}
}
return -1;
}
}
class Node implements Comparable<Node>{
int i;
int j;
int val;
public Node(int i,int j,int val){
this.i = i;
this.j = j;
this.val = val;
}
@Override
public int compareTo(Node n){
return this.val - n.val;
}
}

Time Complexity is O(n²lgn²) since we use PriorityQueue and we need O(lgn²) if we enqueue or dequeue each node into/from a size of n² queue.

Is there any way to find the kth smallest number more efficiently?

Since we can usually do a binary search on a sorted array, is it possible that we do a binary search on the matrix?

Code Logic:

The answer is yes, we will do TWO binary searches in two different parts.

Since the minimum value of the matrix is at matrix[0][0] and the maximum value of the matrix is at matrix[n-1][n-1], we can say that all the numbers are in the range of [matrix[0][0], matrix[n-1][n-1]].
In each iteration, we will keep getting the average number of the starting number and the ending number of the current range (assume the current range is [start, end]), and get the rank of this avg number in the matrix using the helper method:

  1. If the rank of this avg number is less than k, that means the kth number is bigger than this avg number. We will shrink the range to the right, the kth number of the matrix will be in this range: [avg+1, end].
  2. If the rank of this avg number is greater than OR equal to k, that means the kth number is smaller than this avg number. We will shrink the range to the left, the kth number of the matrix will be in this range: [start, avg].

We will stop finding the rank of the avg number until there is only one number in the range.

You might wonder when we using “range” but not the position (i, j) to do the binary search, is it possible that we are passing numbers that are not even in the matrix to the helper method?

The answer is YES, it is highly possible that maybe most of the time, we are trying to get the rank of the numbers that are not in the matrix.

This algorithm will be wrong if we immediately return avg number once we found its rank is k. However, we will still shrink the range to the left when the rank of the avg number is k until there is only one number in the range. i.e., we really found the exact number in the matrix.

👷Let’s see an example to see how we find the 20th smallest number in the matrix by shrinking the range:

I was going to use the above matrix to show the example. But I changed the last two cells to bigger numbers so that we can see how the range varies significantly:

Example Matrix

🔧Step 1. At first, our range is in [1, 1000]. We will call the helper method to get the rank of the avg number 500. Now we know that 500 will be in rank 24 in the matrix. Since the rank of 24 is bigger than our target rank of 21, we will shrink the range to [1, 500].

🔧Step 2. The current range is [1, 500]. We use the helper method to know the rank of its avg number 250 is 24. Again, we shrink the range to [1, 250].

🔧Step 3. The current range is [1, 250]. We use the helper method to know the rank of its avg number 125 is 24. Again, we shrink the range to [1, 125]

🔧Step 4. The current range is [1, 125]. We use the helper method to know the rank of its avg number 63 is 23. Again, we shrink the range to [1, 63].

🔧Step 5. The current range is [1, 63]. We use the helper method to know the rank of its avg number 32 is 16. We found that this time the number is too small, hence the rank of 20 number should be in the range[33, 63].

🔧Step 6. The current range is [33, 63]. We use the helper method to know the rank of its avg number 48 is 22. Again, we shrink the range to [33, 48].

🔧Step 7. The current range is [33, 48]. We use the helper method to know the rank of its avg number 40 is 21. This is the target rank! However, we are not 100% sure if 40 exists in the matrix. We can only sure the 20th number will be in this new range [33, 40].

🔧Step 8. The current range is [33, 40]. We use the helper method to know the rank of its avg number 36 is 18. We found that this time the number is too small, hence the rank of 20 number should be in the range[37, 40].

🔧Step 9. The current range is [37, 40]. We use the helper method to know the rank of its avg number 38 is 18. We found that this time the number is too small, hence the rank of 20 number should be in the range[39, 40].

🔧Step 10. The current range is [39, 40]. We use the helper method to know the rank of its avg number 39 is 19. We found that this time the number is too small, hence the rank of 20 number should be in the range[40, 40].

🚧Currently, there is only one number in the range. Hence, we are 100% sure 40 is the number we want to find and we will return 40 at this point.

After knowing the main code logic, let’s see how we use O(n) to find the rank of the given number in the helper method. It is highly recommended to read Leetcode — 240. Search a 2D Matrix II to know how we can use binary search to find a certain number in a sorted matrix (the matrix condition is the same as this one) in O(2n) when we have an n*n matrix.

After reading Leetcode — 240. Search a 2D Matrix II, you should know why we choose to traverse the whole matrix from the bottom-left. The code logic here is also the same as we did in problem 240 except that this time our goal is not finding the given number, we are going to find how many numbers are less than and equal to the given number.

Since we already mentioned the code logic in the above story, I will just show a simple example of how we count the numbers that met the condition (i.e., ≤target) while we finding the target number in the matrix.

👷Below is an example matrix. We want to see how many numbers are there less than and equal to 20 in the matrix.

Example Matrix

⚠️We use the count variable to represent the number of numbers that are ≤ 20.

🔧Step 1. We found that the current position matrix[4][0] ≤ 20. Since this number is on the bottom row, it means all the numbers in the same column are also less than the target. Also, we will move to the right cell since we want to find a bigger number.
💭count = 5

Step 1

🔧Step 2. We found that the current position matrix[4][1] > 20. We will simply go upward to find a smaller number.
💭count = 5

step 2

🔧Step 3. We found that the current position matrix[3][1] > 20. We will simply go upward to find a smaller number.
💭count = 5

step 3

🔧Step 4. We found that the current position matrix[2][1] > 20. We will simply go upward to find a smaller number.
💭count = 5

step 4

🔧Step 5. We found that the current position matrix[1][1] ≤ 20, which means all the numbers in the same column from this cell and above are also ≤ target. Also, we will move to the right cell since we want to find a bigger number.
💭count = 5 + 2 = 7

step 5

🔧Step 6. We found that the current position matrix[1][2] > 20. We will simply go upward to find a smaller number.
💭count = 7

step 6

🔧Step 7. We found that the current position matrix[0][2] ≤ 20, which means all the numbers in the same column from this cell and above are also ≤ target. Also, we will move to the right cell since we want to find a bigger number.
💭count = 7 + 1 = 8

step 7

🔧Step 8. We found that the current position matrix[0][3] ≤ 20, which means all the numbers in the same column from this cell and above are also ≤ target. Also, we will move to the right cell since we want to find a bigger number.
💭count = 8+ 1 = 9

step 8

🔧Step 9. We found that the current position matrix[0][4] > 20. We will simply go upward to find a smaller number.
💭count = 9

step 9

🚧We have nowhere to go, and we are done finding all the numbers less than equal to 20. We found that there are 9 numbers ≤ 20. If we highlight all these numbers in red, it will be like below —
We can see that they are actually being surrounded by our traversing path! 🆒

result

Below is the 0ms code that beats 100% —

public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
int start = matrix[0][0];
int end = matrix[n-1][n-1];
while(start < end)
{
int mid = start + (end - start)/2;
int rank = helper(matrix, mid);
if(rank < k)
start = mid + 1;
else
end = mid;
}
return start;
}
//count how many numbers <= target
private int helper(int[][] matrix, int target){
int n = matrix.length;
int i = n-1, j = 0, count = 0;
while(i >= 0 && j < n)
{
if(matrix[i][j] <= target)
{
count += (i+1);
j++;
}
else
i--;
}
return count;
}

Time Complexity is O(n*lg(max-min)) where max and min represent the maximum and minimum number in the matrix. We are using binary search on the values in this range, so it will be O(lg(range)) calling the helper method. In the helper method, we will always run O(2*n) since we will either go right or go up in each iteration.

--

--

Anj
Anj

Written by Anj

Click here to view all the stories posted on this blog: https://anj910.medium.com/bookmark-c95e4c52cf44

No responses yet