【leetcode 378】 Kth Smallest Element in a Sorted Matrix（C++）

题目描述：

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.

题目大意：

在一个行和列都有序的矩阵里，找出第k小元素。

思路：

• 用优先队列维护一个大小为k的堆，每次取出堆顶的元素，加入下边的元素（如果是第一行还要加入右边的元素）。操作k次，堆顶元素即为第k小元素。
• 代码：

struct Node {
   int val, i, j;
   Node(int i, int j, int val):i(i), j(j), val(val){}
   bool operator < (const Node & x)const {
   	return val > x.val;
   }
};

class Solution{
public:

   int kthSmallest(vector<vector<int> > &matrix, int k) {
       priority_queue<Node> q;
       int n = matrix.size();
       q.push(Node(0, 0, matrix[0][0]));

       while(--k) {
           Node now = q.top();
           q.pop();
           if (now.i == 0 && now.j + 1 < n) {
               q.push(Node(now.i, now.j+1, matrix[now.i][now.j+1]));
           }
           if (now.i + 1 < n) {
               q.push(Node(now.i+1, now.j, matrix[now.i+1][now.j]));
           }
       }
       return q.top().val;
   }
};

• 二分。最小和最大元素为L和R边界，对mid查找小于它的元素个数，如果小于k，更新L，否则更新R。mid查找过程利用矩阵有序的特性。
• 代码：

class Solution {
public:
   int kthSmallest(vector<vector<int> >& matrix, int k) {
       int n = matrix.size();
       int L = matrix[0][0], R = matrix[n-1][n-1];

       while(L < R){
           //cout << L << "===" << R << endl;
           int mid = (L + R) / 2;
           int temp = search_lower_than_mid(matrix, mid);
           //cout << mid << "---" << temp << endl;
           if (temp < k) L = mid + 1;
           else R = mid;
       }

       return L;
   }

   int search_lower_than_mid(vector<vector<int> >& matrix, int mid) {
       int n = matrix.size();
       int i = n-1, j = 0;
       int cnt = 0;

       while(i >= 0 && j < n) {

           if (matrix[i][j] <= mid) {
               j += 1;
               cnt += i+1;
           }
           else i -= 1;
       }
       return cnt;
   }
};