Kth Smallest Element in a Sorted Matrix
MedPrompt
Given an n x n matrix where each of the rows and columns is sorted in ascending order, return 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 must find a solution with a memory complexity better than O(n2).
Example 1:
Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13Example 2:
Input: matrix = [[-5]], k = 1
Output: -5
Constraints:
n == matrix.length == matrix[i].length1 <= n <= 300-109 <= matrix[i][j] <= 109- All the rows and columns of
matrixare guaranteed to be sorted in non-decreasing order. 1 <= k <= n2
Follow up:
- Could you solve the problem with a constant memory (i.e.,
O(1)memory complexity)? - Could you solve the problem in
O(n)time complexity? The solution may be too advanced for an interview but you may find reading this paper fun.
Approaches
3 approaches with complexity analysis and trade-offs.
A straightforward approach is to treat the matrix as a simple list of numbers. We can iterate through the entire matrix, add all its elements into a single list, and then sort this list. The k-th smallest element will then be the element at the (k-1)-th index of the sorted list.
Algorithm
- Initialize an empty list, say
flatList. - Iterate through each row of then x nmatrix. - For each row, iterate through its elements and add them toflatList. - After iterating through all elements, theflatListwill contain alln^2elements from the matrix. - SortflatListin ascending order. - Return the element at indexk - 1.
Walkthrough
This method ignores the sorted property of the rows and columns during the search, only using it implicitly by the fact that sorting will find the k-th element. The steps are as follows: First, we create a one-dimensional list. Then, we traverse the 2D matrix row by row, and for each element, we add it to our list. This process effectively flattens the matrix into a list of size n*n. Finally, we use a standard sorting algorithm to sort the list and pick the element at index k-1. java import java.util.ArrayList; import java.util.Collections; import java.util.List; class Solution { public int kthSmallest(int[][] matrix, int k) { int n = matrix.length; List<Integer> flatList = new ArrayList<>(); for (int i = 0; i < n; i++) { for (int j = 0; j < n; j++) { flatList.add(matrix[i][j]); } } Collections.sort(flatList); return flatList.get(k - 1); } }
Complexity
Time
O(n^2 * log(n)) - It takes O(n^2) to iterate through the matrix and create the flat list. Sorting a list of n^2 elements takes O(n^2 * log(n^2)), which simplifies to O(n^2 * log(n)).
Space
O(n^2) - We need to create a new list to store all n^2 elements from the matrix.
Trade-offs
Pros
Simple to understand and implement.
Cons
Highly inefficient in both time and space.
The space complexity of O(n^2) violates the problem's memory constraints.
Since each row of the matrix is sorted, we can view this problem as finding the k-th smallest element from n sorted lists. A min-heap is a perfect data structure for this task. We can start by adding the first element of each row to the min-heap. Then, we repeatedly extract the minimum element from the heap and add the next element from the same row back into the heap. After k extractions, the last element extracted is our answer.
Algorithm
- Create a min-heap to store elements in the format
[value, rowIndex, colIndex]. The heap will be ordered byvalue. - Add the first element of each of thenrows to the min-heap. - Loopktimes: - Extract the minimum element[value, row, col]from the heap. Thisvalueis the next smallest element. - If the extracted element's row has more elements (i.e.,col + 1 < n), add the next element(matrix[row][col + 1], row, col + 1)to the heap. - The last value extracted afterkiterations is the answer.
Walkthrough
This approach leverages the fact that each row is a sorted list. We are essentially merging n sorted lists. A min-heap (or Priority Queue) is used to efficiently find the minimum among the current heads of all n lists. We initialize the heap with the first element of each row. Then, we extract the minimum element k times. Each time we extract an element, we add the next element from its row into the heap. This ensures the heap always holds the smallest candidates for the next overall smallest element. java import java.util.PriorityQueue; class Solution { public int kthSmallest(int[][] matrix, int k) { int n = matrix.length; PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> a[0] - b[0]); for (int i = 0; i < n; i++) { minHeap.offer(new int[]{matrix[i][0], i, 0}); } int result = -1; for (int i = 0; i < k; i++) { int[] current = minHeap.poll(); result = current[0]; int row = current[1]; int col = current[2]; if (col + 1 < n) { minHeap.offer(new int[]{matrix[row][col + 1], row, col + 1}); } } return result; } }
Complexity
Time
O(k * log(n)) - Initializing the heap with `n` elements takes O(n * log(n)). We then perform `k` extractions, each followed by a potential insertion. Both operations take O(log(n)) time. The total complexity is dominated by the `k` operations.
Space
O(n) - The min-heap will store at most `n` elements at any given time (one from each row).
Trade-offs
Pros
Much more memory-efficient than the brute-force approach, meeting the problem's constraints.
It's a standard and effective pattern for merging k sorted lists.
Cons
The time complexity is dependent on
kand can be slow ifkis large (e.g., close to n^2).
This is the most optimal approach. The k-th smallest element must be within the range of values present in the matrix, from the smallest (matrix[0][0]) to the largest (matrix[n-1][n-1]). We can perform a binary search on this range of values to find the answer. For each mid value we pick during the binary search, we need to efficiently count how many elements in the matrix are less than or equal to mid.
Algorithm
- Define a search range for the answer:
low = matrix[0][0]andhigh = matrix[n-1][n-1]. - Whilelow <= high: - Pick amidvalue in the range. - Count how many elements in the matrix are less than or equal tomid. This can be done in O(n) time. - If thecountis less thank, it means the answer must be larger thanmid, so we setlow = mid + 1. - If thecountis greater than or equal tok,midcould be the answer, or the answer is smaller. We sethigh = mid - 1to search for a potentially smaller answer. - The loop terminates whenlow > high, andlowwill be the k-th smallest element.
Walkthrough
The binary search is not on the indices, but on the value range of the matrix elements. For any chosen value mid, we can determine the number of elements in the matrix that are less than or equal to it. Let's call this count(mid). If count(mid) < k, our guess mid is too small, and we must search for a larger value. If count(mid) >= k, our guess mid is a potential answer, but there might be a smaller value that also satisfies this, so we try to find a smaller one. The key is the efficient O(n) counting function. It works by starting at the bottom-left corner (row=n-1, col=0). If matrix[row][col] <= mid, we know all row+1 elements in that column are also <= mid, so we add row+1 to the count and move to the next column. Otherwise, the element is too large, and we move up a row. java class Solution { public int kthSmallest(int[][] matrix, int k) { int n = matrix.length; int low = matrix[0][0]; int high = matrix[n - 1][n - 1]; while (low <= high) { int mid = low + (high - low) / 2; int count = countLessOrEqual(matrix, mid); if (count < k) { low = mid + 1; } else { high = mid - 1; } } return low; } private int countLessOrEqual(int[][] matrix, int x) { int n = matrix.length; int count = 0; int row = n - 1; int col = 0; while (row >= 0 && col < n) { if (matrix[row][col] <= x) { count += (row + 1); col++; } else { row--; } } return count; } }
Complexity
Time
O(n * log(D)) - Where D is the difference between the maximum and minimum values in the matrix. The binary search performs log(D) iterations, and each iteration involves an O(n) counting step.
Space
O(1) - This approach uses only a few variables to keep track of the search range and count, requiring constant extra space.
Trade-offs
Pros
Extremely efficient in both time and space.
The O(1) space complexity is optimal.
The time complexity is generally the best possible for this problem.
Cons
The logic, especially the O(n) counting step and the binary search boundary conditions, can be less intuitive to grasp compared to the heap approach.
Solutions
Solution
import java.util.Comparator ; import java.util.PriorityQueue ; public class Kth_Smallest_Element_in_a_Sorted_Matrix { class Node { int x ; int y ; int val ; public Node ( int x , int y , int val ) { this . x = x ; this . y = y ; this . val = val ; } } class Solution { public int kthSmallest ( int [][] matrix , int k ) { if ( matrix == null || matrix . length == 0 || k <= 0 ) { return 0 ; } PriorityQueue < Node > pq = new PriorityQueue <>( new Comparator < Node >() { @Override public int compare ( Node o1 , Node o2 ) { return o1 . val - o2 . val ; } }); for ( int i = 0 ; i < matrix . length ; i ++) { pq . offer ( new Node ( i , 0 , matrix [ i ][ 0 ])); // 第一列入heap } while (! pq . isEmpty ()) { Node n = pq . poll (); if ( k == 1 ) { return n . val ; } if ( n . y + 1 < matrix . length ) { pq . offer ( new Node ( n . x , n . y + 1 , matrix [ n . x ][ n . y + 1 ])); // 入current的row的下一个 } k --; } return 0 ; } } class Solution_optimize { public int kthSmallest ( int [][] matrix , int k ) { if ( matrix == null || matrix . length == 0 || k <= 0 ) { return 0 ; } int n = matrix . length ; int min = matrix [ 0 ][ 0 ]; int max = matrix [ n - 1 ][ n - 1 ]; while ( min < max ) { int mid = min + ( max - min ) / 2 ; int smallerCount = countSmallerThanMid ( matrix , mid ); if ( smallerCount < k ) { min = mid + 1 ; } else { max = mid ; } } return min ; // return max will also work } private int countSmallerThanMid ( int [][] matrix , int mid ) { int count = 0 ; // start from top-right corner, where all its left is smaller but all its below is larger int i = 0 ; int j = matrix [ 0 ]. length - 1 ; while ( i < matrix [ 0 ]. length && j >= 0 ) { if ( matrix [ i ][ j ] > mid ) { j --; } else { count += j + 1 ; i ++; } } return count ; } } } ////// class Solution { public int kthSmallest ( int [][] matrix , int k ) { int n = matrix . length ; int left = matrix [ 0 ][ 0 ], right = matrix [ n - 1 ][ n - 1 ]; while ( left < right ) { int mid = ( left + right ) >>> 1 ; if ( check ( matrix , mid , k , n )) { right = mid ; } else { left = mid + 1 ; } } return left ; } private boolean check ( int [][] matrix , int mid , int k , int n ) { int count = 0 ; int i = n - 1 , j = 0 ; while ( i >= 0 && j < n ) { if ( matrix [ i ][ j ] <= mid ) { count += ( i + 1 ); ++ j ; } else { -- i ; } } return count >= k ; } }''' time complexity of the countSmallerThanMid(mid) function is O(n), as it iterates through at most n rows and columns of the matrix. number of iterations in the binary search `while left < right:` is O(log(maxVal - minVal)), where maxVal - minVal represents the range of values in the matrix. overall time complexity of the code is `O(n * log(maxVal - minVal))` no extra space ''' class Solution : def kthSmallest ( self , matrix : List [ List [ int ]], k : int ) -> int : def countSmallerThanMid ( mid ): count = 0 # start from top-right corner, where all its left is smaller but all its below is larger i , j = 0 , n - 1 while i < n and j >= 0 : if matrix [ i ][ j ] > mid : j -= 1 else : count += j + 1 # j is index, j+1 is count for that row i += 1 return count n = len ( matrix ) left , right = matrix [ 0 ][ 0 ], matrix [ n - 1 ][ n - 1 ] while left < right : mid = ( left + right ) >> 1 if countSmallerThanMid ( mid ) >= k : # The kth smallest element is smaller than or equal to mid, so update the right boundary right = mid else : # The kth smallest element is larger than mid, so update the left boundary left = mid + 1 # If the loop terminates, left and right will be equal # Return left (or right), which represents the kth smallest element return left # also passing OJ: return right ############ ''' This algorithm runs in O(n) time complexity because each every element is added to the heap only once and is visited only once The space complexity is also O(n) because we store at most n elements in the heap and at most n elements in the visited set ''' import heapq class Solution ( object ): def kthSmallest ( self , matrix , k ): """ :type matrix: List[List[int]] :type k: int :rtype: int """ visited = {( 0 , 0 )} heap = [( matrix [ 0 ][ 0 ], ( 0 , 0 ))] while heap : # checking every element val , ( i , j ) = heapq . heappop ( heap ) k -= 1 if k == 0 : return val if i + 1 < len ( matrix ) and ( i + 1 , j ) not in visited : heapq . heappush ( heap , ( matrix [ i + 1 ][ j ], ( i + 1 , j ))) visited . add (( i + 1 , j )) if j + 1 < len ( matrix [ 0 ]) and ( i , j + 1 ) not in visited : heapq . heappush ( heap , ( matrix [ i ][ j + 1 ], ( i , j + 1 ))) visited . add (( i , j + 1 ))// OJ: https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/ // Time: O(Klog(N)) // Space: O(N) class Solution { typedef tuple < int , int , int > Item ; public: int kthSmallest ( vector < vector < int >>& A , int k ) { int N = A . size (), dirs [ 2 ][ 2 ] = { { 0 , 1 },{ 1 , 0 } }; vector < vector < bool >> seen ( N , vector < bool > ( N )); priority_queue < Item , vector < Item > , greater <>> pq ; pq . emplace ( A [ 0 ][ 0 ], 0 , 0 ); seen [ 0 ][ 0 ] = true ; while ( -- k ) { auto [ n , x , y ] = pq . top (); pq . pop (); for ( auto & [ dx , dy ] : dirs ) { int a = x + dx , b = y + dy ; if ( a < 0 || a >= N || b < 0 || b >= N || seen [ a ][ b ]) continue ; seen [ a ][ b ] = true ; pq . emplace ( A [ a ][ b ], a , b ); } } return get < 0 > ( pq . top ()); } };Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.