[leetcode] 378. Kth Smallest Element in a Sorted Matrix

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note: 
You may assume k is always valid, 1 ≤ k ≤ n2.

这道题是在行列均排好序的矩阵中找第k小的数字，题目难度为Medium。

最直观的想法是用大根堆来存储当前最小的k个数字，然后遍历矩阵，如果遍历位置数字大于等于堆顶元素，跳过该行继续遍历，否则将数字存入堆中并删除堆顶数字（保证堆中有k个数字）。最终堆顶数字即是第k小的数字。具体代码：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        priority_queue<int> heap;
        for(int i=0; i<matrix.size(); ++i) {
            for(int j=0; j<matrix[0].size(); ++j) {
                if(heap.size() < k) heap.push(matrix[i][j]);
                else {
                    if(heap.top() <= matrix[i][j]) break;
                    else {
                        heap.push(matrix[i][j]);
                        heap.pop();
                    }
                }
            }
        }
        return heap.top();
    }
};

还可以逐次取出最小的数字来找出最终结果。将第一行元素存入小根堆，最小元素必定在堆顶（matrix[0][0]），然后删除堆顶数字并用它同列的下一数字代替，这样当前最小的数字依然在堆顶，遍历k次后即可得到第k小的数字。由于需要记录数字的坐标，这里用结构体包含了数字值和两个坐标。具体代码：

class Solution {
    struct num {
        int val;
        int x, y;
        num(int v, int xx, int yy) : val(v), x(xx), y(yy) {}
    };
    struct cmp {
        bool operator() (const num& n1, const num& n2) {
            return n1.val > n2.val;
        }
    };
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        priority_queue<num, vector<num>, cmp> heap;
        int sz = matrix.size();
        for(int i=0; i<sz; ++i) 
    ​    ​    ​heap.push(num(matrix[0][i], 0, i));
        for(int i=0; i<k-1; ++i) {
            num curMin = heap.top();
            heap.pop();
            if(curMin.x < sz-1) {
                heap.push(num(matrix[curMin.x+1][curMin.y], curMin.x+1, curMin.y));
            }
        }
        return heap.top().val;
    }
};

看了别人的代码，还有用二分查找来处理题目的，分享给大家。在二分查找循环中，统计矩阵中小于等于中间值的数字个数，拿它和k比较来确定第k小的数字在左半部分还是右半部分。具体代码：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        int bgn = matrix[0][0], end = matrix[n-1][n-1];
        while(bgn < end) {
            int mid = (bgn + end) / 2;
            int cnt = 0;
            for(int i=0; i<n; ++i) {
                cnt += (upper_bound(matrix[i].begin(), matrix[i].end(), mid)-matrix[i].begin());
            }
            if(cnt < k) bgn = mid + 1;
            else end = mid;
        }
        return bgn;
    }
};