Scale Engineer

Kth Smallest Element in a Sorted Matrix

Med
#0365Time: 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.4 companies
Algorithms
Binary SearchSorting
Data structures
ArrayHeap (Priority Queue)Matrix
Companies
SalesforcePhonePeXOKX

Prompt

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 13

Example 2:

Input: matrix = [[-5]], k = 1
Output: -5

 

Constraints:

  • n == matrix.length == matrix[i].length
  • 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

 

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 the n x n matrix. - For each row, iterate through its elements and add them to flatList. - After iterating through all elements, the flatList will contain all n^2 elements from the matrix. - Sort flatList in ascending order. - Return the element at index k - 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.

Solutions

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

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.

← Previous

Previous problem

Next →

Next problem