Kth Smallest Element in a Sorted Matrix

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.

How to Solve

Algorithm

• Start with left = minOfMatrix = matrix[0][0] and right = maxOfMatrix = matrix[n-1][n-1].
• Find the mid of the left and the right. This middle number is NOT necessarily an element in the matrix.
• If countLessOrEqual(mid) >= k, we keep current ans = mid and try to find smaller value by searching on the left side. Otherwise, we search on the right side.
• Since ans is the smallest value which countLessOrEqual(ans) >= k, so it’s the k^th smallest element in the matrix.

How to count number of elements less or equal to x efficiently?

• Since our matrix is sorted in ascending order by rows and columns.
• We use two pointers, one points to the rightmost column c = n-1, and one points to the lowest row r = 0.
  • If matrix[r][c] <= x then the number of elements in row r less or equal to x is (c+1) (Because row[r] is sorted in ascending order, so if matrix[r][c] <= x then matrix[r][c-1] is also <= x). Then we go to next row to continue counting.
  • Else if matrix[r][c] > x, we decrease column c until matrix[r][c] <= x (Because column is sorted in ascending order, so if matrix[r][c] > x then matrix[r+1][c] is also > x).
• Time complexity for counting: O(M+N).

Test Cases

Input:

(int[]) matrix = [[1,5,9],[10,11,13],[12,13,15]]
(int) k = 8

Output:

(int) 13

Explanation:

The elements in the matrix are [1,5,9,10,11,12,13,13,15], 
and the 8th smallest number is 13

Solution

C++ Java Python

class Solution {
public:
    int m, n;
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        m = matrix.size(), n = matrix[0].size();
        int left = matrix[0][0], right = matrix[m-1][n-1], ans = -1;
        while (left <= right) {
            int mid = (left + right) >> 1;
            if (countLessOrEqual(matrix, mid) >= k) {
                ans = mid;
                right = mid - 1;
            } else left = mid + 1;
        }
        return ans;
    }
    int countLessOrEqual(vector<vector<int>>& matrix, int x) {
        int cnt = 0, c = n - 1;
        for (int r = 0; r < m; ++r) {
            while (c >= 0 && matrix[r][c] > x) --c;
            cnt += (c + 1);
        }
        return cnt;
    }
};

class Solution {
    int m, n;
    public int kthSmallest(int[][] matrix, int k) {
        m = matrix.length; n = matrix[0].length;
        int left = matrix[0][0], right = matrix[m-1][n-1], ans = -1;
        while (left <= right) {
            int mid = (left + right) >> 1;
            if (countLessOrEqual(matrix, mid) >= k) {
                ans = mid;
                right = mid - 1;
            } else left = mid + 1;
        }
        return ans;
    }
    int countLessOrEqual(int[][] matrix, int x) {
        int cnt = 0, c = n - 1;
        for (int r = 0; r < m; ++r) {
            while (c >= 0 && matrix[r][c] > x) --c;
            cnt += (c + 1);
        }
        return cnt;
    }
}

class Solution:
    def kthSmallest(self, matrix, k):
        m, n = len(matrix), len(matrix[0])

        def countLessOrEqual(x):
            cnt = 0
            c = n - 1
            for r in range(m):
                while c >= 0 and matrix[r][c] > x:
                    c -= 1
                cnt += (c + 1)
            return cnt

        left, right = matrix[0][0], matrix[-1][-1]
        ans = -1
        while left <= right:
            mid = (left + right) // 2
            if countLessOrEqual(mid) >= k:
                ans = mid
                right = mid - 1
            else:
                left = mid + 1

        return ans