378. Kth Smallest Element in a Sorted Matrix 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 ≤ n². 378. Kth Smallest Element in a Sorted Matrix

---

给定一个矩阵，该矩阵每行和每列都是有序的，要找出第k小的数。 如果只利用每行或者每列是有序的，则可以使用堆解决。我们以行为例，那么每行都是有序的，则相当于把所有行进行n路归并，然后找到第k小的数。所以我们构建一个最小堆，首先把所有行的首元素放到堆中，然后n路归并k次，每次弹出堆顶元素（即当前最小值），然后把原堆顶元素所在行的下一个元素放入堆中。k次归并结束之后，堆顶元素即为第k小的元素。 代码如下：

class Solution {
private:
    struct Node {
        int val, row, col;
        Node(int v = 0, int r = 0, int c = 0)
            : val(v)
            , row(r)
            , col(c){};
        friend bool operator<(const Node& n1, const Node& n2) { return n1.val > n2.val; }
    };
public:
    int kthSmallest(vector<vector<int> >& matrix, int k)
    {
        int n = matrix.size();
        priority_queue<Node> pn;
        for (int i = 0; i < n; ++i)
            pn.push(Node(matrix[i][0], i, 0));
        Node t;
        while (k–) {
            t = pn.top();
            pn.pop();
            if (t.col < n – 1)
                pn.push(Node(matrix[t.row][t.col + 1], t.row, t.col + 1));
        }
        return t.val;
    }
};

本代码提交AC，用时45MS。时间复杂度是O(klgn)，堆的大小为n，每次归并调整堆需要lgn时间，一共需要归并k次，所以总的时间复杂度是O(klgn)。
另外用二分的方法可以同时利用行有序和列有序的特点。这种矩阵可以看成一个曲面，曲面的最低点为左上角元素left，最高点为右下角元素right，所以曲面就是从左上角向上延伸到右下角（想象一下MATLAB画的图）。那么第k小的数肯定在[left,right]范围内，我们可以二分查找这个区间。假设中点是mid，则我们在每一行找找比mid小的数有多少个，加起来如果等于cnt，则说明mid这个数排在第cnt位。因为每一行也是有序的，所以可以直接用upper_bound查找第一个比mid大的数。
代码如下：

class Solution {
public:
    int kthSmallest(vector<vector<int> >& matrix, int k)
    {
        int n = matrix.size(), left = matrix[0][0], right = matrix[n – 1][n – 1];
        while (left <= right) {
            int mid = left + (right – left) / 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)
                left = mid + 1;
            else
                right = mid – 1;
        }
        return left;
    }
};

本代码提交AC，用时43MS。时间复杂度是O(lgX nlgn)，X为[left,right]的区间大小，ngln是对于每个mid，都需要在每一行二分查找比mid小的数的个数。
还可以将上述复杂度降为O(lgX n)。就是在找cnt，不需要每一行都二分查找，我们可以从矩阵的左下角往右上角阶梯查找，如果要查找的数更大，则往右移，否则往上移，每次查询只需要O(n)的时间，所以总时间是O(lgX n)。
代码如下：

class Solution {
private:
    int count_less_equal(vector<vector<int> >& matrix, int num)
    {
        int n = matrix.size(), i = n – 1, j = 0, cnt = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= num) {
                cnt += i + 1;
                ++j;
            }
            else
                –i;
        }
        return cnt;
    }
public:
    int kthSmallest(vector<vector<int> >& matrix, int k)
    {
        int n = matrix.size(), left = matrix[0][0], right = matrix[n – 1][n – 1];
        while (left <= right) {
            int mid = left + (right – left) / 2;
            int cnt = count_less_equal(matrix, mid);
            if (cnt < k)
                left = mid + 1;
            else
                right = mid – 1;
        }
        return left;
    }
};

本代码提交AC，用时39MS。