378. Kth Smallest Element in a Sorted Matrix

1. Description

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

5. Solutions

Binary Search

n = matrix.size()
Time complexity: O(nlog(max(matrix) - min(matrix)))
Space complexity: O(1)

// both row and column are ascending
// while merge sort only use row's feature
class Solution {
    int count_smaller_numbers(const vector<vector<int>> &matrix, int number) {
        const int rows = matrix.size();
        const int columns = matrix.front().size();

        int result = 0;
        for (int i = rows - 1, j = 0; i >= 0 && j < columns;) {
            matrix[i][j] < number ? result += (i + 1), ++j : --i;
        }

        return result;
    }

public:
    int kthSmallest(const vector<vector<int>> &matrix, int k) {
        int result = matrix.front().front();
        for (int left = matrix.front().front(), right = matrix.back().back(); left < right;) {
            int mid = (left + right) / 2;

            count_smaller_numbers(matrix, mid) < k ? result = mid, left = mid + 1 : right = mid;
        }

        return result;
    }
};