【堆】B005_有序矩阵中第 K 小的元素（穷举+最小堆 | 最大堆 | 二分+穷举 | 双二分 | 最优解）

一、题目描述

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.

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

return 13.

二、题解

方法一：穷举+最小堆

先将所有元素入堆，最后从堆内取出第 $k$ 个元素返回即可。

/**
 * @date: 2/20/2020 9:27 PM
 * @Execution info：39ms，9.7%，MB，5.13%
 */
public int kthSmallest(int[][] matrix, int k) {

  PriorityQueue<Integer> pQ = new PriorityQueue<>();
  for (int[] row : matrix)
  for (int i : row)
    pQ.add(i);
    
  while (--k > 0) {
    pQ.poll();
  }
  return pQ.peek();
}

复杂度分析

• 时间复杂度： $O(N^2 × logN)$，
• 空间复杂度： $O(N)$，

---

方法二：最大堆（优化）

思路

• 建立最小堆，当堆里面的元素个数大于 $k$ 个时，将堆顶元素抛出。
• 由于数组是有序的，所以队列中的第 $k$ 一定是最大的；也就说，达到 k 个以后，顺序添加矩阵的任何数值的数都可保证最后堆顶元素就是第 $k$ 小的数。

/**
 * @date: 2/20/2020 9:37 PM
 * @Execution info：23ms，34.61%，MB，5.13%
 */
public int kthSmallest1(int[][] matrix, int k) {

  PriorityQueue<Integer> pQ = new PriorityQueue<>((e1, e2) -> e2.compareTo(e1));
  for (int[] row : matrix)
  for (int i : row) {
    pQ.add(i);
    if (pQ.size() > k)
      pQ.poll();
  }
  return pQ.peek();
}

复杂度分析

• 时间复杂度： $O(N^2 × logk)$，
• 空间复杂度： $O(k)$，

---

方法三：二分 + 穷举

思路与算法

• 找出矩阵中最大、最小值的中间值 mid=(left+right) / 2；在二维矩阵中寻找小于等于 mid 的元素个数 count。
• 若 count < k，表明第 k 小的数在范围的右边且不包含 mid，搜索区域的左边界变为：left=mid+1。
• 若 count >= k，表明第 k 小的数在范围的左边且可能包含 mid，搜索区域的右边界变为：right = mid。

/**
 * @date: 2/20/2020 10:00 PM
 * @Execution info：4ms，62.83%，MB，5.13%
 */
public int kthSmallest2(int[][] matrix, int k) {

  final int M = matrix.length;
  final int N = matrix[0].length;
  int l = matrix[0][0];
  int r = matrix[M-1][N-1];

  while (l < r) {
    int mid = l + ((r - l) >>> 1);
    int count = 0;
    count = search(matrix, mid);
    if (count < k)
      l = mid + 1;
    else
      r = mid;
  }
  return l;
}

//在矩阵中寻找，小于mid的数的个数
private int search(int[][] M, int mid) {
  int count = 0;
  for (int[] row : M)
  for (int i : row)
    if (i <= mid)
      count++;
  return count;
}

复杂度分析

• 时间复杂度： $O(N^2 × log\ D)$， $D$ 为矩阵中最大值与最小值之差。
• 空间复杂度： $O(1)$，

---

方法四：双二分（查下标）

既然二维矩阵的每一行都是有序的，那么可以通过查找每一行第一个大于 target 的元素的索引，进而确定二维矩阵中每一行大于 target 的个数 count。

/**
 * @date: 2/20/2020 10:56 PM
 * @Execution info：2ms，64.65%，MB，5.13%
 */
public int kthSmallest3(int[][] matrix, int k) {

  final int M = matrix.length;
  final int N = matrix[0].length;
  int l = matrix[0][0];
  int r = matrix[M-1][N-1];

  while (l < r) {
    int mid = l + ((r - l) >>> 1);
    int count = 0;
    for (int[] row : matrix) {
      count += bSearch(row, mid);
    }
    if (count < k)
      l = mid + 1;
    else
      r = mid;
  }
  return l;
}
// 查找第一个大于 target 的数的下标。
private int bSearch(int[] arr, int target) {

  int l = 0, r = arr.length;
  while (l < r) {
    int mid = l + ((r - l) >>> 1);
    if (arr[mid] > target)
      r = mid;
    else
      l = mid + 1;
  }
  return l;
}

复杂度分析

• 时间复杂度： $O( N × log N × log D)$， $D$ 为矩阵中最大值与最小值之差。
• 空间复杂度： $O(1)$，

---

方法五：特解（没想出来）

该算法的目标也是寻找小于小于 target 的数的个数。

想法

• 由于二维数组的每一列都是有序的，一个特别的方法就是从数组的左下角 $[N-1][0]$ 开始查找，
• 如果比 target 小，我们就向右移一位，为什么不上移，因为我们知道当前列的当前位置的上方的所有的数字都小于 target。 cnt += i+1。
• 如果比 target 小/等于，则向上移一位，这样我们也能算出 cnt 的值。
• 其余部分跟上面的方法相同。

复杂度分析

• 时间复杂度： $O()$，
• 空间复杂度： $O()$，