Kth Smallest Element in a Sorted Matrix#

Problem statement#

[1]You are given an n x n matrix where each row and column is sorted in ascending order. Your task is to find the k-th smallest element in this matrix.

Please note that we are looking for the k-th smallest element based on its position in the sorted order, and not counting distinct elements.

Additionally, it is required to find a solution with a memory complexity better than O(n^2).

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.
-10^9 <= matrix[i][j] <= 10^9.
All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
1 <= k <= n^2.

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.

Solution 1: Transform the 2-D matrix into a 1-D vector then sort#

You can implement exactly what Example 1 has explained.

Code#

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
int kthSmallest(const vector<vector<int>>& matrix, int k) {
    vector<int> m;
    // transform the 2D matrix into a 1D array m
    for (auto& row : matrix) {
        m.insert(m.end(), row.begin(), row.end());
    }
    // sort the array m
    sort(m.begin(), m.end());
    return m.at(k - 1);
}
int main() {
    vector<vector<int>> matrix{{1,5,9},{10,11,13},{12,13,15}};
    cout << kthSmallest(matrix, 8) << endl;
    matrix = {{-5}};
    cout << kthSmallest(matrix, 1) << endl;
}

Output:
13
-5

The core idea behind this solution is to transform the 2D matrix into a 1D sorted array, making it easier to find the k-th smallest element efficiently. The time complexity of this solution is dominated by the sorting step, which is O(N*logN), where N is the total number of elements in the matrix.

Complexity#

Runtime: O(N*logN), where N = n^2 is the total number of elements in the matrix.
Extra space: O(N).

Solution 2: Build the max heap and keep it ungrown#

Instead of sorting after building the vector in Solution 1, you can do the other way around. It means building up the vector from scratch and keeping it sorted.

Since you need only the k-th smallest element, std::priority_queue[2] can be used for this purpose.

Code#

#include <iostream>
#include <vector>
#include <queue>
using namespace std;
int kthSmallest(const vector<vector<int>>& matrix, int k) {
    priority_queue<int> q;
    for (int row = 0; row < matrix.size(); row++) {
        for (int col = 0; col < matrix[row].size(); col++) {
            q.push(matrix[row][col]);
            // maintain q's size does not exceed k
            if (q.size() > k) {
                q.pop();
            }
        }
    }
    return q.top();
}
int main() {
    vector<vector<int>> matrix{{1,5,9},{10,11,13},{12,13,15}};
    cout << kthSmallest(matrix, 8) << endl;
    matrix = {{-5}};
    cout << kthSmallest(matrix, 1) << endl;
}

Output:
13
-5

Complexity#

Runtime: O(N*logk), where N = n^2 is the total number of elements of the matrix.
Extra space: O(k).

Conclusion#

Solution 2 maintains a priority queue of size k, allowing it to efficiently keep track of the k-th smallest element encountered while iterating through the matrix.

This approach is handy for large matrices, as it doesn’t require sorting the entire matrix.

Exercise#

Find K Pairs with Smallest Sums[3].