Finding the Kth Smallest Element _present_ in a Sorted Matrix

Question

My question is inspired by this particular SO comment which is unanswered (I am facing that problem myself):

I am trying to find the Kth smallest element in a sorted matrix:

Given an n x n matrix where each of the rows and columns are sorted in ascending order, return the kth smallest element in the matrix.
Note that it is the kth smallest element in the sorted order, not the kth distinct element.

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.

The code I have is this:

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int lowest=INT_MAX, highest=INT_MIN;
        for(vector<int> row: matrix) {
            lowest=min(lowest, row[0]);
            highest=max(highest, row[row.size()-1]);
        }
        
        while(lowest<highest) {
            int mid=lowest+(highest-lowest+1)/2;
            int places=0;
            
            for(int i=0; i<matrix.size(); i++) {
                places+=(upper_bound(matrix[i].begin(), matrix[i].end(), mid)-matrix[i].begin());
            }
            if(places<=k) lowest=mid;      //this gives a number _not_ in the matrix (14)
            else highest=mid-1;
            // if(places<k) lowest=mid+1;  //this gives a number _in_ the matrix (13)
            // else highest=mid;   //also use mid=lowest+(highest-lowest)/2 instead;
        }
        
        return lowest;
    }
};

The commented code returns a number that is present in the matrix (13), whereas the uncommented code returns 14, which is not in the matrix.

What gives? What is the intuition behind finding the number present in the matrix? Working code on ideone here.

Answer 1

Let me give you another approach to solve that problem, since the rows and columns are sorted ascending, for every element elem in the matrix at position (x, y), the element at (x + 1, y) and the element (x, y + 1) are both smaller than elem, if those were valid indexes.

Then we know for sure that the smallest element in the matrix is matrix[0][0], and that if we remove that element, the smallest element would be the smallest between matrix[0][1] and matrix[1][0]. And so on if we keep removing elements.

Our algorithm is going to put in a min heap (priority queue) every element that is candidate at some point for being the smallest element in the matrix remaining (we will not modify the matrix directly), at first the heap will only contain the element at (0, 0) and it indexes in the matrix so we can locate it quickly, represented as pair<int, pair<int, int>>.

Our steps will be the following, pop the smallest element in the heap, put the two new candidates for being the smallest in the heap, and repeat until we pop k elements.

You might have noticed that there will be repeated numbers in our heap because a number can be inserted in the heap by the element at its top and the element at its left, and we dont want that, so lets solve that issue.

Now we are only getting the element at the right if the element we pop is in the first row. If we do this there will not be any repeated indexes in the heap, but, will this work ?, the answer to that question is yes, since the columns and rows are sorted, it is impossible that we get a minimun without getting all elements that are over it and at its left, and we are still making sure that there is at least a way to get to every element in ascending order.

You should also notice that we at most insert 2*k elements in the heap.

Implementation details: I assume there are at least k elements in the matrix.

typedef pair<int, int> pii;
typedef pair<int, pair<int, int>> matrixValue;

int kthSmallest(vector<vector<int>>& matrix, int k) {
    priority_queue<matrixValue, vector<matrixValue>, greater<>> priorityQueue;

    matrixValue smallest = matrixValue(matrix[0][0], pii(0, 0));

    priorityQueue.push(smallest);

    for (int i = 1; i < k; ++i) {
        matrixValue iterSmallest = priorityQueue.top();
        priorityQueue.pop();

        int x = iterSmallest.second.first;
        int y = iterSmallest.second.second;

        if ( x == 0 and y + 1 < matrix[0].size() ) {
            priorityQueue.push(matrixValue(matrix[x][y + 1], pii(x, y + 1)));
        }
        if ( x + 1 < matrix.size() ) {
            priorityQueue.push(matrixValue(matrix[x + 1][y], pii(x + 1, y)));
        }
    }

    return priorityQueue.top().first;
}