Matrix: Kth smallest element in a sorted matrix

Problem Statement:

You are given n*n matrix, where every row and column is sorted in non descending order.
You need to find kth smallest element in the given 2D array.

Example

Input:
k = 3  
array =
        10, 20, 30, 40
        50, 60, 70, 80
        90, 100, 110, 120
        130, 140, 150,160

Output: 30

Explanation: The 3rd smallest element is 20 

Solution

The solution is very simple. We need to use modified binary search to solve the problem.

But to use Binary Search, we dont have a sorted array, but we do have a matrix.

SO here we go on the number range instead of index range.

WKT, the smallest number is at the top left corner and highest number at the bottom right corner.

There two will represent the range. We can use binary search as below:

Staring Index = mat[0][0]
Ending Index = mat[n-1][n-1]

Steps:

1. Find the middle of start and end

2. Count the number smaller than or equal to middle.
While counting, keep track of the “smallest number greater than the middle (n1)” and at the same time “biggest number less than or equal to middle (n2)”

3. If the count is equal to “k”, n1 is our required number as it is the “biggest number less than or equal to middle”

4. If the count is less than “k”, we can update start = n2 to search the higher part of the matrix and if the count is greater than “k”, then we can update end = n1 to search in lower part of the matrix.

Example:

Here K = 5

Solution in C++

#include <vector>    
#include <algorithm>  
//visit www.ProDeveloperTutorial.com for 450+ solved questions  
#include <iostream>    
#include <string>
#include <unordered_set>

using namespace std;


int countLessEqual(vector<vector<int>> &matrix, int mid, pair<int, int> &smallLargePair) 
{
    int count = 0;
    int n = matrix.size(), row = n - 1, col = 0;
    while (row >= 0 && col < n) 
    {
      if (matrix[row][col] > mid) 
      {
        smallLargePair.second = min(smallLargePair.second, matrix[row][col]);
        row--;
      } else {
        smallLargePair.first = max(smallLargePair.first, matrix[row][col]);
        count += row + 1;
        col++;
      }
    }
    return count;
}



int get_Kth_smallest_element(vector<vector<int>> &matrix, int k) 
{
    int n = matrix.size();
    int start = matrix[0][0];
    int end = matrix[n - 1][n - 1];

    while (start < end) 
    {
      int mid = start + (end - start) / 2;

      pair<int, int> smallLargePair = {matrix[0][0], matrix[n - 1][n - 1]};
      int count = countLessEqual(matrix, mid, smallLargePair);
      
      if (count == k) 
      {
        return smallLargePair.first;
      }

      if (count < k) {
        start = smallLargePair.second;  // search higher
      } else {
        end = smallLargePair.first;  // search lower
      }
    }
    return start;
  }

int main()
{
   vector<vector<int>> mat = {
        { 10, 20, 30, 40 },
        { 15, 25, 35, 45 },
        { 25, 29, 37, 48 },
        { 32, 33, 39, 50 },
    };
    cout << "7th smallest element is "
         << get_Kth_smallest_element(mat, 7);
    return 0;
}

Output:

7th smallest element is 30