Kth Smallest Element in Sorted 2D Matrix Problem

Given a 2D matrix mat of integers with n rows and m columns. Each row of the matrix contains m integers in ascending order and also each column contains n integers in ascending order. Your task is to find the kth smallest element in matrix mat. Your algorithm should be efficient enough, so that it suffice q such queries.

Input Format:

First and second parameter of the function kth_smallest_element, is n and m, depicting the number of rows and columns in the matrix. The third parameter is the integer k, denoting the kth smallest element to be returned by the function. The fourth and last parameter of the function is the 2D matrix mat.

Output Format:

Return an integer denoting the kth smallest element of the input 2D matrix.

Constraints:

-10^9

Sample Test Case:

Sample Input:

n = 3

m = 3

mat = [ [1,2,3] , [4,5,6] , [7,8,9] ]

q = 1

k = 6

Sample Output:

Explanation:

There are total 9 elements in the matrix, which when arranged in ascending order are as follows: [1,2,3,4,5,6,7,8,9]. Now, the 6th smallest element is none other than 6 itself and hence the answer.

Solution

We have provided solutions which contain necessary comments to understand the approach used:

1) brute_force_solution:

Description:

The brute-force approach for this problem is pretty much simple. We can just push all the elements of the matrix mat into a linear container like an array. Then our task it to find the kth smallest element in the resulting array. We can solve this easily by first sorting the resulting array and hence then the (k-1)th index element will be the kth smallest element (assuming 0-based indexing).

Time Complexity:

O(n*m*log(n*m)) where n is number of rows of matrix and m is number of columns of matrix.

Sorting the array takes O(size*log(size)) time (where size = n*m) and accessing the (k-1)th element is a O(1) time operation. Hence, the total Time Complexity of the solution will be O(n*m*log(n*m)).

Auxiliary Space Used:

O(n*m) where n is number of rows of matrix and m is number of columns of matrix.

To store array containing all elements of matrix, and as size of matrix will be n*m.

Space Complexity:

O(n*m) where n is number of rows of matrix and m is number of columns of matrix.

For storing input it will take space of O(n*m) and auxiliary space used is O(n*m).

So, O(n*m) + O(n*m) → O(n*m).


// -------- START --------

// @param mat: 2D matrix with rows and column sorted in ascending order
// @param k: integer k, denoting kth element to be returned
int kth_smallest_element(int n,int m, int k,vector < vector > &mat)
{   
    // container to store all the elemnts of the 2D matrix
    vector all_elements;

    // storing all elements of 2D matrix in a linear array
    for (int i = 0; i < n; i++)
    {
        for (int j = 0; j < m; j++)
        {
            all_elements.push_back(mat[i][j]);
        }
    }

    // sorting the linear array container
    sort(all_elements.begin(), all_elements.end());

    // returning kth element
    return all_elements[k - 1];
}

// -------- END --------

‍

2) other_solution:

Description:

BST based solution. Here, we will be maintaining a BST of size n (number of rows in mat).

We will initialise the BST by adding elements of the first column of the matrix mat. Now, we will start extracting the min-elements from the BST. Each time we extract an element from the BST we insert another element in the BST which will be the next element in the same row as of extracted element. Suppose the element extracted is x and it belongs to i row and j column. So, we will insert the next smallest element after x from the ith row. But as according to given constraint the columns are also in sorted order. Hence, the next smallest element will be the (j+1)th element. So, we insert the element that belongs to ith row and (j+1)th column of matrix in the BST. Doing so ensures that every time we extract the minimum element from the matrix.

So, we repeat this extract and insert operation for k times and hence, the last element extracted in this sequence will be the kth smallest element in the matrix mat. We will use hashmap to map elements in the heap to its row and column number so that we can easily point to the next element to be inserted.

Time Complexity:

O((n+k)*log(n)) where n is number of rows of matrix and k is the element number which need to be extracted.

To build a BST by adding first element of each row will take O(n*log(n)), extracting elements k times and inserting elements k time will take O(k*log(n)).

So, time complexity will be O(n*log(n)) + O(k*log(n)) → O((n+k)*log(n)).

Auxiliary Space Used:

O(n) where n is number of rows of matrix.

We are maintaining a BST of size O(n) and a hashmap of size O(n) to store the row and columns of the elements in the BST.

Space Complexity:

O(n*m) where n is number of rows of matrix and m is number of columns of matrix.

Input space complexity of O(n*m), the size of the given matrix mat and auxiliary space used id O(n).

So, O(n*m) + O(n) → O(n*m)


// -------- START --------

// @param mat: 2D matrix with rows and column sorted in ascending order
// @param k: integer k, denoting kth element to be returned
int kth_smallest_element(int n, int m, int k, vector> &mat)
{
    // min-heap
    set> min_heap;

    // iterator to point to the extract min_heap element
    set>::iterator it;

    // hash-map to store current idx of every row
    // maps row_index to same_row's_element_index
    unordered_map idx_hash_map;

    // initializing min-heap and hash map
    for (int i = 0; i < n; i++)
    {
        // inserting the first elemnt of each row
        min_heap.insert(make_pair(mat[i][0], i));
        // mapping element index for each row.
        // As for each row we are inserting their 0th positon
        // element. So, mapping to 0th idx for each row
        idx_hash_map[i] = 0;
    }
    int result = -1;
    while (k--)
    {
        // getting minimum element from the heap
        it = min_heap.begin();
        result = it->first;

        // getting the row of the minimum element extracted
        int curr_element_row_num = it->second;

        // getting the column idx of the extracted element in the
        // current element row
        int pos_in_curr_row = idx_hash_map[curr_element_row_num];

        // removing the extracted minimum elemnt from the min heap
        min_heap.erase(it);

        // if the extracted element is not the last one in its row
        if (pos_in_curr_row != (m - 1))
        {
            // go to next smallest number in that row
            pos_in_curr_row++;
            // insert the next smallest number in the min heap
            min_heap.insert(make_pair(mat[curr_element_row_num][pos_in_curr_row],
                                      curr_element_row_num));
            // update the hash map after the insertion
            idx_hash_map[curr_element_row_num] = pos_in_curr_row;
        }
    }
    // returning the final result
    return result;
}

// -------- END --------

3) optimal_solution:

Description:

In this solution, instead of directly figuring out the kth smallest element we search the input space of integers i.e from -10^9 to 10^9 to check if it is the kth smallest number in matrix. We will feed value X from -10^9 to 10^9 to getRankOfX() function and that tells us the number of elements smaller than and equal to X. The smallest integer X, for which the getRankOfX() function evaluates to k, is the kth smallest element in the matrix and hence, the solution.

Now let’s examine the behaviour of the getRankOfX() function. It seems that the function getRankOfX() is a monotonically increasing function(i.e. as we increase the value of X, the function evaluates to a higher value).

So, we can use binary search to find the correct value of X. Now, let’s say we evaluated the getRankOfX() function on some integer X and the function evaluates to a value val which is greater than k. Then obviously the kth smallest element will be less than X and if the val was less than k, then in that case the kth smallest element will be greater than current X.

We use the above discussed guidelines to search the kth smallest element by reducing the search space on every iteration. Also, one point worth noticing is that we can reduce our search space to the range [smallest_element_in_mat, largest_element_in_mat] i.e. in range [top-leftmost element, bottom-rightmost element].

So, till now everything is pretty simple. Now, only one black box is pending – “getRankOfX()” function. It is fairly easy to evaluate the getRankOfX() function as the matrix is sorted both row-wise and column-wise. We start moving from the top-rightmost cell and keep track of all the elements greater than X. If the value of the current cell is greater than X, than that means that all the values below that cell in the current column will be greater than X and in this case we move a column left. If the value of the cell is less than X, then we move down a cell in the same column. We continue this until we either reach the last row or first column.

Once, these iterations ends we will have a count of all the elements greater than X. We subtract it from total elements in the matrix and that is the evaluated value for getRankOfX() for a particular X. Kindly, refer to the solution for better understanding.

Time Complexity:

The time complexity of the getRankOfX() function is O(m+n) as in worst case it travels the manhattan distance from top-rightmost cell to bottom-left most cell ( i.e it makes m row traversals and n column traversals). Now, the getRankOfX() function is called log(R) times because of binary search on the range R, where R is the search range (max_element_of_mat – min_element_of_mat).

Hence, the total Time Complexity of this solution is O(log(R)*(m+n)), where R is the search space ( bottom right most element minus the top left most element).

Auxiliary Space Used:

It’s worth noticing that we did not used any extra memory for any of our calculation and hence, the auxiliary space complexity is O(1).

Space Complexity:

O(n*m) where n is number of rows of matrix and m is number of columns of matrix.

Input space complexity of O(n*m), the size of the given matrix mat and auxiliary space used id O(1).

So, O(n*m) + O(1) → O(n*m)


// -------- START --------
// @param mat: 2D matrix with rows and column sorted in ascending order
// @param X : Element whose rank is to be determinded in mat
// @returns : number of elements less than equal to X in mat
int getRankOfX(int X, vector> &mat) {
    int cols = mat[0].size();
    int rows = mat.size();
    int x = 0, y = cols - 1;
    // stores elements greater than X
    int cnt = 0;
    // we start from top-right most cell
    while (y >= 0 and x < rows)
    {   
        // if current cell val is > X
        // adding all the values below it to cnt
        // as all of them will be greater than X
        // Then,moving a column to left
        if (mat[x][y] > X) {
            cnt += rows - x;
            y--;
        }
        // if current cell val is <= X
        // moving a cell down in same column
        else {
            x++;
        }
    }
    // Elements less than equal to X = Total Elements - cnt
    return (rows * cols - cnt);
}

// @param mat: 2D matrix with rows and column sorted in ascending sorted
// @param k: integer k, denoting kth element to be returned
int kth_smallest_element(int n, int m, int k, vector> &mat)
{   
    // lower bound of search
    int lo = mat[0][0];
    // upper bound of search
    int hi = mat[n-1][m-1];
    // maximum number of iterations to reach the kth
    // smallest element
    int iterations = log2(hi-lo) + 1;
    int mid;
    for (int i = 0; i < iterations; i++) {
        // evaluating mid as our current answer
        mid = (hi+lo) >> 1;
        // if rankOfX is greater than k
        // then our answer lies left to mid
        if(getRankOfX(mid,mat) >=k)
            hi = mid;
        else
        // if rankOfX is less than k
        // then our answer lies strictly right to mid
            lo = mid+1;
    }
    return lo;
}
// -------- END --------