问题描述
给定一个 n x n 的矩阵，其中每行和每列元素均按升序排序，找出矩阵中第 k 小的元素。
示例：
输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
解释：矩阵中的元素为 [1,5,9,10,11,13,12,13,15]，第 8 小元素是 13
算法实现
1. 二分查找法
实现思路
- 利用矩阵的有序性，采用二分查找的思路
- 确定搜索范围：矩阵左上角元素（最小值）和右下角元素（最大值）
- 对于每个中间值 mid，统计矩阵中小于等于 mid 的元素个数
- 根据统计结果调整搜索范围，直到找到第 k 小的元素
关键实现
public int kthSmallest(int[][] matrix, int k) {
int left = matrix[0][0];
int right = matrix[matrix.length - 1][matrix[0].length - 1];
while (left <= right) {
int mid = (left + right) / 2;
int temp = count(matrix, mid);
if (temp >= k) {
right = mid - 1;
} else {
left = mid + 1;
}
}
return left;
}
计数方法 (count)
- 从矩阵左下角（或右上角）开始搜索
- 对于当前值，如果小于等于 mid，则整行左边的元素都符合条件，count 加上该行符合条件的元素数
- 如果大于 mid，则向上移动一行继续搜索
private int count(int[][] matrix, int mid) {
int count = 0;
int j = matrix[0].length - 1;
for (int[] ints : matrix) {
while (j >= 0 && ints[j] > mid) {
j--;
}
if (j < 0) {
break;
}
count += (j + 1);
}
return count;
}
复杂度分析
- 时间复杂度：O(n log(max - min))，其中 n 是矩阵的维度，max 和 min 分别是矩阵中的最大值和最小值
- 二分查找的次数为 log(max - min)，每次查找需要 O(n) 的时间来统计元素个数
- 空间复杂度：O(1)，只使用了常数级别的额外空间
2. 优先队列法
实现思路
- 利用最小堆进行多路归并
- 首先将矩阵的第一列元素加入堆中，每个元素记录其值、行号和列号
- 然后执行 k 次操作：每次取出堆顶元素（当前最小元素），并将其所在行的下一个元素加入堆中
- 第 k 次取出的元素即为矩阵中第 k 小的元素
关键实现
public int kthSmallest(int[][] matrix, int k) {
Queue<Integer[]> queue = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
for (int i = 0; i < matrix[0].length; i++) {
queue.add(new Integer[]{matrix[i][0], i, 0});
}
int result = 0;
for (int i = 0; i < k; i++) {
Integer[] topArr = queue.poll();
result = topArr[0];
if (topArr[2] < matrix[topArr[1]].length - 1) {
queue.add(new Integer[]{matrix[topArr[1]][topArr[2] + 1], topArr[1], topArr[2] + 1});
}
}
return result;
}
复杂度分析
- 时间复杂度：O(k log n)，其中 n 是矩阵的维度
- 堆的大小最多为 n，每次堆操作的时间复杂度为 O(log n)
- 总共需要执行 k 次堆操作
- 空间复杂度：O(n)，堆中最多存储 n 个元素
3. 简单多路归并法
实现思路
- 维护一个索引数组，记录每一行当前比较到的列索引
- 每次从所有行的当前元素中找出最小值，记录该值并将对应行的索引加 1
- 重复 k 次上述操作，第 k 次找到的最小值即为矩阵中第 k 小的元素
关键实现
public int kthSmallest(int[][] matrix, int k) {
int[] indexArr = new int[matrix.length];
int res = 0;
for (int i = 0; i < k; i++) {
int min = Integer.MAX_VALUE;
int ji = 0;
for (int j = 0; j < matrix.length; j++) {
if (indexArr[j] >= matrix[0].length) {
continue;
}
if (min > matrix[j][indexArr[j]]) {
min = matrix[j][indexArr[j]];
ji = j;
}
}
indexArr[ji]++;
res = min;
}
return res;
}
复杂度分析
- 时间复杂度：O(k * n)，其中 n 是矩阵的维度
- 每次需要遍历 n 行来找出当前最小值
- 总共需要执行 k 次这样的遍历
- 空间复杂度：O(n)，需要一个大小为 n 的索引数组
算法比较
| 算法 | 时间复杂度 | 空间复杂度 | 特点 |
|---|---|---|---|
| 二分查找法 | O(n log(max - min)) | O(1) | 高效，尤其当 max - min 较小时性能最优 |
| 优先队列法 | O(k log n) | O(n) | 适用于 k 较小的情况，当 k 接近 n² 时性能不如二分查找法 |
| 简单多路归并法 | O(k * n) | O(n) | 实现简单，但时间复杂度较高，适用于小矩阵或 k 较小的情况 |
适用场景
-
二分查找法：
- 适用于大多数情况，尤其是矩阵较大或 k 接近 n² 时
- 空间效率高，只需要常数级别的额外空间
-
优先队列法：
- 适用于 k 较小的情况，如寻找前几个最小元素
- 实现相对直观，利用了堆的特性
-
简单多路归并法：
- 适用于小矩阵或作为教学演示
- 实现简单，易于理解，但性能较差
输入输出示例
// 测试用例 1
输入：matrix = {{5}}, k = 1
输出：5
// 测试用例 2
输入：matrix = {{1, 3}, {2, 4}}, k = 2
输出：2
// 测试用例 3
输入：matrix = {{1, 5, 9}, {10, 11, 13}, {12, 13, 15}}, k = 8
输出：13
// 测试用例 4
输入：matrix = {{1, 5, 9}, {10, 11, 13}, {12, 13, 15}}, k = 5
输出：11
