Finding kth smallest element in a sorted matrix

Problem Statement

Given a sorted matrix of size n*n and a number k. Find the k^th smallest element in the matrix. The matrix can contain duplicates.
A sorted matrix is a matrix in which each row and each column is sorted.

Example Test Case

Given Matrix:
1 2 3
3 5 9
7 7 10
K: 4
Expected Output: 3
Explanation: Since the 4th smallest element among all the elements of this matrix is 3.

Solution

Since the matrix is sorted, the smallest element of the whole matrix is located at position (0, 0) and the largest element is located at position (n-1, n-1). Let us call the element located at (0, 0) as start and the element located at (n-1, n-1) as end.

It is clear that the desired answer exists b/w start and end. Let us call our answer answer. To find answer, we can apply binary search b/w start and end.

Let us assume that we have a function called findNumSmallerThan which takes a number target and tells the total count of elements in the matrix which are smaller than or equal to target.

When we apply binary search b/w start and end, we will count the no of numbers smaller than or equal to mid using the above defined function findNumSmallerThan. Let us call the value returned by findNumSmallerThan as val. Now two cases arise:

1. val < k: In this case, the answer will be greater than mid and we will set lo = mid + 1 in our binary search.
2. val >= k: In this case, one of the possible answers could be mid (if mid exists in the matrix). If mid doesn’t exist in the matrix then the answer is some number smaller than mid. Since, we are not sure which of these subcases apply, we will keep on going down in our search space and set hi = mid - 1

findNumSmallerThan(target)

This function returns the number of elements smaller than or equal to target. Let us call this number val.
To find val, we can sum up the number of elements smaller than or equal to target for each row of the matrix. Since each row of the matrix is itself sorted, I can find out the number of elements smaller than or equal to target by using another binary search.
Since C++ already provides such function called upper_bound, I will directly use that as shown below:

Implementation

Time Complexity

Since, we are doing binary search b/w start and end, which are smallest and largest element in the matrix respectively and for each binary search operation, we are further doing n binary searches (within each row) for calculating val. The total time complexity would be log(end - start)*n*log(n)