378. Kth Smallest Element in a Sorted Matrix

2023, Jan 15 2 mins read

1. Description

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the $k^{th}$ smallest element in the matrix.
Note that it is the $k^{th}$ smallest element in the sorted order, not the $k^{th}$ distinct element.
You must find a solution with a memory complexity better than O($n^2$).

2. Example

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

3. 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$

4. Solutions

Binary Search

m = matrix.size(), n = matrix.front().size()
Time complexity: O((m+n)log(max(matrix) - min(matrix)))
Space complexity: O(1)

class Solution {
public:
    int kthSmallest(const vector<vector<int>> &matrix, int k) {
        const int m = matrix.size(), n = matrix.front().size();
        int left = matrix[0][0], right = matrix[m - 1][n - 1];
        while (left < right) {
            int middle = left + (right - left) / 2;

            if (count_equal_smaller_numbers(matrix, middle) < k) {
                left = middle + 1;
            } else {
                right = middle;
            }
        }

        return left;
    }

private:
    int count_equal_smaller_numbers(const vector<vector<int>> &matrix, int middle) {
        const int m = matrix.size(), n = matrix.front().size();
        int count = 0, i = m - 1, j = 0;
        while (i >= 0 && j < n) {

            if (matrix[i][j] <= middle) {
                count += i + 1;
                ++j;
            } else {
                --i;
            }
        }

        return count;
    }
};