【Leetcode】378. Kth Smallest Element in a Sorted Matrix

题目地址：

https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/

给定一个二维矩阵$A$，矩阵每行从左到右单调递增，每列从上到下单调递增。求其中的第$k$小元素。相同元素如果出现多次则算多次。

二分答案。显然整个矩阵的最小值就是$A$的最左上角元素，而最大值就是$A$的最右下角元素。所以我们可以以这个为区间，二分寻找满足这样条件的数，这个数是使得$A$中小于等于它的数的个数大于等于$k$的最小数。很显然，这样的数就是$A$中的第$k$小的数。而要计算$A$中小于等于$c$的数的个数，我们可以从矩阵右上角出发，如果遇到小于等于$c$的数，说明其自己和其左方都是满足条件的数，计入计数，并向下走一步；如果遇到大于$c$的数，说明其自己和其右方都不满足条件，直接向左一步。直到走到左下角为止。代码如下：

class Solution {
 public:
  int kthSmallest(vector<vector<int>>& g, int k) {
    int n = g.size();
    int l = g[0][0], r = g[n - 1][n - 1];
    auto cnt = [&](int z) {
      int x = 0, y = n - 1, res = 0;
      while (y >= 0 && x < n) {
        if (g[x][y] <= z)
          res += y + 1, x++;
        else
          y--;
      }
      return res;
    };
    while (l < r) {
      int mid = l + (r - l >> 1);
      if (cnt(mid) >= k)
        r = mid;
      else
        l = mid + 1;
    }

    return l;
  }
};

时间复杂度$O(n\log r)$，其中$r$是$A$中的数据范围。空间$O(1)$。