LeetCode 378. 有序矩阵中第 K 小的元素 Java 实现

题目分析

给定一个 n x n 矩阵，其中每行和每列都按升序排列，找出矩阵中第 k 小的元素（k 从 1 开始计数）。

由于矩阵行列均有序，可以使用二分查找值域的方法：

· 矩阵最小值在左上角 matrix[0][0]，最大值在右下角 matrix[n-1][n-1]。
· 在 [min, max] 范围内二分查找一个数 mid，统计矩阵中 ≤ mid 的元素个数 count。
· 如果 count < k，说明第 k 小的元素比 mid 大，令 left = mid + 1。
· 否则，第 k 小的元素 ≤ mid，令 right = mid。
· 最终 left 即为答案。

统计 ≤ mid 的元素个数时，利用行列有序的特性，从右上角开始：

· 若 matrix[row][col] <= mid，则当前行从第 0 列到第 col 列的所有元素都 ≤ mid，故 count += col + 1，然后 row++。
· 否则，col–。
· 这样可以在 O(n) 时间内完成统计。

Java 实现

classSolution{publicintkthSmallest(int[][]matrix,intk){intn=matrix.length;intleft=matrix[0][0];intright=matrix[n-1][n-1];while(left<right){intmid=left+(right-left)/2;intcount=countLessOrEqual(matrix,mid);if(count<k){left=mid+1;}else{right=mid;}}returnleft;}// 统计矩阵中 <= target 的元素个数privateintcountLessOrEqual(int[][]matrix,inttarget){intn=matrix.length;introw=0;intcol=n-1;intcount=0;while(row<n&&col>=0){if(matrix[row][col]<=target){count+=col+1;// 这一行从 0 到 col 都 <= targetrow++;}else{col--;}}returncount;}}

代码详解

· 二分范围：left = matrix[0][0]，right = matrix[n-1][n-1]，即矩阵的最小值和最大值。
· 统计函数：从右上角 (0, n-1) 出发，利用行升序和列升序的性质，每次移动一行或一列，时间复杂度 O(n)。
· 二分判断：count < k 说明 mid 太小，答案在右半部分；否则答案在左半部分（包括 mid）。
· 循环结束时 left == right，即为第 k 小的元素。

样例追踪

以 matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8 为例：

· 初始 left = 1, right = 15。
· mid = 8，统计 ≤ 8 的个数：1,5 → 2 个。2 < 8，left = 9。
· mid = 12，统计 ≤ 12 的个数：1,5,9,10,11,12 → 6 个。6 < 8，left = 13。
· mid = 13，统计 ≤ 13 的个数：1,5,9,10,11,12,13,13 → 8 个。8 >= 8，right = 13。
· left == right == 13，返回 13。

结果正确。

复杂度分析

· 时间复杂度：O(n log(max - min))，其中 n 为矩阵边长，max 和 min 为矩阵中的最大值和最小值。每次二分 O(log(max-min))，每次统计 O(n)。
· 空间复杂度：O(1)，仅使用常数个变量。

扩展：最小堆解法

也可以使用最小堆归并 n 个有序数组：

· 将每行的第一个元素入堆。
· 弹出 k-1 次，每次弹出后将该元素所在行的下一个元素入堆。
· 第 k 次弹出的元素即为答案。

classSolution{publicintkthSmallest(int[][]matrix,intk){intn=matrix.length;PriorityQueue<int[]>pq=newPriorityQueue<>((a,b)->a[0]-b[0]);for(inti=0;i<n;i++){pq.offer(newint[]{matrix[i][0],i,0});}for(inti=0;i<k-1;i++){int[]cur=pq.poll();introw=cur[1],col=cur[2];if(col+1<n){pq.offer(newint[]{matrix[row][col+1],row,col+1});}}returnpq.peek()[0];}}

· 时间复杂度：O(k log n)，适合 k 较小的情况。
· 空间复杂度：O(n)。

推荐：二分法更通用，不受 k 大小影响，是本题的最优解之一。