378. 有序矩阵中第 K 小的元素

前言

本文解法着重于用$O(1)$的空间复杂度解决，而时间复杂度为$O(nlog(r-l))$。另外的解法还有转化成多个有序列表合并，或者用最小堆做的。

二分查找思路

题目 https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/
官方题解 https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/solution/you-xu-ju-zhen-zhong-di-kxiao-de-yuan-su-by-leetco/

利用二分查找，去找到那个第k大的数字。首先设置左右区间，left为左上角，right为右下角。mid为两者平均值。

假设取mid=8，从左下角开始，沿着蛇形路线找出所有不大于mid的数。计算这样的数有几个(6+6+4+2=18)。

然后，比较它与k。如果中间值取太大了，则应当寻找左半区间，否则应当寻找右半区间。

图例

以下图为例，矩阵中有[…, 1, 3, 6, 10]这些数，第k个数是3。将矩阵中的数字有序铺开于下图。刚开始，left和right分处图两端。取left和right中间的值mid，统计≤mid的元素数rank，并比较rank与k的关系：

• 当mid取到红色区间时，rank<k，会促使分界线向右移动。
• 否则，矩阵取到蓝色区间时，rank≥k，会促使分界线向左移动。
  • 留意到，当搜索范围的中位值取到元素3和6之间，即3、4或5时，rank的值都是k，但只有当取3时，目标值才找到。
    

假如第k小的数字是3，当搜索范围收敛到该值时，答案已确定。

代码

代码如下：

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int left = matrix[0][0], right = matrix[n-1][n-1];
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (binarySearch(matrix, k, mid)) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }

    public boolean binarySearch(int[][] matrix, int k, int mid) {
        int n = matrix.length;
        int i = n-1, j = 0;
        int rank = 0;
        while (i >= 0 && j <= n-1) {
            if (matrix[i][j] <= mid) {
                rank += i + 1;
                j++;
            } else {
                i--;
            }
        }
        return rank >= k;
    }
}

需要注意的细节有两个。
第一个细节是设置能向左收缩的中位值。计算中位值时使用：

 int mid = left + (right - left) / 2;

而不使用

int mid = (left + right) / 2;

因为，尽管后者在计算正数时，能保证mid < right，比如当left和right取[3, 4]时，mid=3。但在计算负数时，会取到右侧，比如[-4, -3]时，mid=-3。
而前者不论计算正数还是负数，都能保证mid<right，从而令区间能向左侧收缩。

第二细节是搜索区间应鼓励收缩。区间最终要收敛到一个点，即只要满足rank>=k，区间都应该向左收缩，直到收敛到一个点为止。否则如上图，3 4 5都能满足rank>=k都能取到，选择4或者5就错了。