在有序矩阵中找到第 K 个最小元素 _present_

Question

我的问题受到this particular SO comment 的启发，没有答案（我自己也面临这个问题）：

我是trying to findK排序矩阵中最小的元素：

给定一个 n x n 矩阵，其中每一行和每一列都按升序排序，返回矩阵中第 k 个最小的元素。
请注意，它是排序顺序中第 k 个最小的元素，而不是第 k 个不同的元素。

输入：矩阵 = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
说明：矩阵中的元素为[1,5,9,10,11,12,13,13,15]，第8小数为13。

我的代码是这样的：

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;
    }
};

注释的代码返回矩阵中存在的数字 (13)，而未注释的代码返回矩阵中不的14。

什么给了？找到矩阵中存在的数字背后的直觉是什么？ ideone 上的工作代码here。

Answer 1

让我给你另一种方法来解决这个问题，因为行和列是按升序排序的，对于矩阵中位置 (x, y) 处的每个元素 elem，(x + 1, y) 处的元素并且元素 (x, y + 1) 都小于elem，如果它们是有效索引的话。

那么我们肯定知道矩阵中最小的元素是matrix[0][0]，如果我们去掉那个元素，最小的元素就是matrix[0][1]和matrix[1][0之间的最小值]。如果我们继续删除元素，依此类推。

我们的算法将把每个在某个时刻候选成为矩阵中剩余最小元素的元素放入一个最小堆（优先队列）中（我们不会直接修改矩阵），起初堆只会包含 (0, 0) 处的元素，它在矩阵中建立索引，因此我们可以快速定位它，表示为 pair>。

我们的步骤如下，弹出堆中最小的元素，将两个新的候选元素放入堆中，然后重复直到弹出k元素。

您可能已经注意到我们的堆中会有重复的数字，因为可以通过顶部元素和左侧元素将数字插入堆中，我们不希望这样，所以让我们解决这个问题。

现在，如果我们弹出的元素在第一行，我们只会获取右边的元素。如果我们这样做，堆中不会有任何重复的索引，但是，这会工作吗？这个问题的答案是肯定的，因为列和行是排序的，所以我们不可能在没有得到所有元素的情况下得到一个最小值在它的上方和左侧，我们仍在确保至少有一种方法可以按升序访问每个元素。

您还应该注意到，我们在堆中最多插入 2*k 个元素。

实现细节： 我假设矩阵中至少有 k 个元素。

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;
}