code378- 有序矩阵中第 K 小的元素
实现思路
1 方法1：暴力解法
- S1 将矩阵中的所有元素放入一个数组中
- S2 对数组进行排序
- S3 返回数组中第 k 小的元素
- 时间复杂度：O(n^2 * log(n))
- 空间复杂度：O(n^2)
题目特性：
- 矩阵的有序性：每行和每列都是升序的
2 方法2：值域二分法
- S1 转换题意： 第k小的数 等价于 所有 <=某个数x 的数量cnt，cnt >= k个条件下的 最小值
- 即 Min(v1, v2, v3 ...), 其中 vn 都满足 Cnt(<=vn值的数量) >= k
- 其实 画个值域图 就能直观理解了
- S2 开区间二分查找：保证不会出现mid死循环的情况
3 方法3：多路归并/ 最小堆
- S1 初始化堆
- S2 把第一列的元素放入堆
- S3 弹出k-1次，每次弹出后把该行的下一列元素加入堆
- S4 返回堆顶元素，此时即第k小的元素
参考文档
代码实现
1 方法1: 二分查找
-
时间复杂度：O(nlogU)
- 其中 n 是 matrix 的行数和列数
- U = matrix[n−1][n−1] − matrix[0][0]
- 二分 O(logU) 次, 每次需要跑一个 O(n) 的双指针
-
空间复杂度： O(1)
function kthSmallest(matrix: number[][], k: number): number {
const n = matrix.length;
// 转换题意： 第k小的数 等价于 所有 <=某个数x 的数量cnt，cnt >= k个条件下的 最小值
// 即 Min(v1, v2, v3 ...), 其中 vn 都满足 Cnt(<=vn值的数量) >= k
// 其实 画个值域图 就能直观理解了
let [l, r] = [matrix[0][0] - 1, matrix.at(-1).at(-1)];
// 开区间二分查找：
// 这种写法保证了 区间内只剩两个数时就退出循环，即 取 [min-1, r]
// 不会出现 left 和 right 相等，造成 mid死循环的情况
while (l + 1 < r) {
const mid = l + ((r - l) >> 1);
//即 Cnt(<=mid值的 数量) >= k，说明此时 该值过大了，右边界需要缩短
if (check(mid)) {
r = mid;
} else {
// 说明此时 该值过小了，左边界需要扩大
l = mid;
}
}
return r;
// Util: 检查<=目标值的 元素个数是否达到k个
function check(target: number): boolean {
let [row, col, cnt] = [0, n - 1, 0];
// 每次都从 右上角开始遍历
while (row < n && col >= 0 && cnt < k) {
// 当前列的值太大，左移一列
if (matrix[row][col] > target) {
col--;
} else {
// 当前行从 [0,col]都小于等于target
cnt += col + 1;
row++;
}
}
return cnt >= k;
}
}
2 方法2：多路归并/ 最小堆
- 时间复杂度：O(klogn)
- 归并 k 次，每次堆中插入和弹出的操作时间复杂度均为 logn
- 注意，k 在最坏情况下是 n^2
- 空间复杂度：O(n)，堆的大小始终为 n
type IItem = { num: number; pos: number[] };
function kthSmallest(matrix: number[][], k: number): number {
// S1 初始化堆
const n = matrix.length;
const heap = minHeap<IItem>((a, b) => a.num < b.num);
// S2 把第一列的元素放入堆
for (let row = 0; row < n; row++) {
heap.add({ num: matrix[row][0], pos: [row, 0] });
}
// S3 弹出k-1次，每次弹出后把该行的下一列元素加入堆
for (let i = 0; i < k - 1; i++) {
const [row, col] = heap.extract().pos;
if (col < n - 1) {
heap.add({ num: matrix[row][col + 1], pos: [row, col + 1] });
}
}
// S4 返回堆顶元素，此时即第k小的元素
return heap.peek().num;
}
function minHeap<T extends IItem>(compare: (a: T, b: T) => boolean) {
const heap: T[] = [];
return {
size: () => heap.length,
isEmpty: () => heap.length === 0,
peek: () => heap[0],
add: (item: T) => {
heap.push(item);
siftUp(heap.length - 1);
},
extract: () => {
const ret = heap[0];
swap(0, heap.length - 1);
heap.pop();
siftDown(0);
return ret;
},
};
function siftUp(idx: number) {
while (idx > 0) {
const pdx = ~~((idx - 1) / 2);
// 说明此时 cur < parent，需要上浮
const willUp = pdx >= 0 && compare(heap[idx], heap[pdx]);
if (!willUp) break;
swap(idx, pdx);
idx = pdx;
}
}
function siftDown(idx: number) {
while (1) {
let ldx = idx * 2 + 1,
rdx = ldx + 1;
let ndx = idx;
if (ldx < heap.length && compare(heap[ldx], heap[ndx])) ndx = ldx;
if (rdx < heap.length && compare(heap[rdx], heap[ndx])) ndx = rdx;
if (ndx === idx) break;
swap(idx, ndx);
idx = ndx;
}
}
function swap(i: number, j: number) {
[heap[i], heap[j]] = [heap[j], heap[i]];
}
}
