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

🟠   🔖  数组 二分查找 矩阵 排序 堆（优先队列）  🔗 力扣 LeetCode

题目

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order , not the kth distinct element.

You must find a solution with a memory complexity better than O(n^2).

Example 1:

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

Output: 13

Explanation: The elements in the matrix are [1,5,9,10,11,12,13, _ 13_ ,15], and the 8th smallest number is 13

Example 2:

Input: matrix = [[-5]], k = 1

Output: -5

Constraints:

• n == matrix.length == matrix[i].length
• 1 <= n <= 300
• -10^9 <= matrix[i][j] <= 10^9
• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
• 1 <= k <= n^2

Follow up:

• Could you solve the problem with a constant memory (i.e., O(1) memory complexity)?
• Could you solve the problem in O(n) time complexity? The solution may be too advanced for an interview but you may find reading this paper fun.

题目大意

给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。 请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。

你必须找到一个内存复杂度优于 O(n^2) 的解决方案。

解题思路

思路一：二分查找

由于每一行和每一列都是升序排序的，可以尝试对目标值进行二分查找，然后计算小于等于目标值的元素个数。具体步骤如下：

1. 设置二分查找的左边界 left 为矩阵中最小的元素，右边界 right 为矩阵中最大的元素。
2. 在每一次迭代中，计算中间值 mid，然后遍历整个矩阵，统计小于等于 mid 的元素个数。
3. 如果小于等于 mid 的元素个数大于等于 k，说明目标值在左半部分，更新 right = mid。
4. 如果小于等于 mid 的元素个数小于 k，说明目标值在右半部分，更新 left = mid + 1。
5. 重复上述步骤直到 left 和 right 相遇，最终 left 的值就是目标值。

复杂度分析

• 时间复杂度： O(n log(max-min))，其中 max 和 min 分别是矩阵中的最大值和最小值。因为进行二分查找的次数是 log(max-min) 级别的，而每次二分查找的时间复杂度是 O(n)。
• 空间复杂度： O(1)。

---

思路二：堆（优先队列）

可以使用最小堆（Min Heap）实现一个大小为 k 的优先队列。最小堆是一个二叉树，其中每个节点的值都小于或等于其子节点的值。

1. 将矩阵的第一列元素全部插入优先队列中。这时队列中的元素个数是矩阵的行数 n。
2. 进行以下循环操作：
   • 弹出队首元素，这是当前队列中最小的元素。
   • 如果该元素所在行还有下一个元素，将下一个元素插入队列。
   • 重复上述两步操作直到找到第 k 小的元素。
3. 当完成 k 次循环后，队首元素即为第 k 小的元素。

这样每次都弹出队列中最小的元素，保证队列中保留着当前阶段的最小的 k 个元素。最终， k 次循环后，队首元素即为矩阵中第 k 小的元素。

复杂度分析

• 时间复杂度：O(k * log(n))，其中 n 是矩阵的边长。由于每次插入和删除的操作都需要 log(n) 的时间，循环执行 k - 1 次，因此总时间复杂度是 O(k * log(n))。
• 空间复杂度：O(n)，空间复杂度主要取决于优先队列的大小，最坏情况下为 O(n)。

代码

二分查找

/**
 * @param {number[][]} matrix
 * @param {number} k
 * @return {number}
 */
var kthSmallest = function (matrix, k) {
	const n = matrix.length;
	let left = matrix[0][0];
	let right = matrix[n - 1][n - 1];

	while (left < right) {
		const mid = Math.floor((left + right) / 2);
		let count = 0;
		let j = n - 1;
		for (let i = 0; i < n; i++) {
			while (j >= 0 && matrix[i][j] > mid) {
				j--;
			}
			count += j + 1;
		}
		if (count >= k) {
			right = mid;
		} else {
			left = mid + 1;
		}
	}
	return left;
};

堆（优先队列）

/**
 * @param {number[][]} matrix
 * @param {number} k
 * @return {number}
 */
var kthSmallest = function (matrix, k) {
	let heap = [];
	const add = ([val, i, j]) => {
		heap.push([val, i, j]);
		heapifyUp(heap.length - 1);
	};
	const pop = () => {
		if (heap.length == 0) {
			return null;
		}
		const min = heap[0];
		const max = heap.pop();
		if (heap.length > 0) {
			heap[0] = max;
			heapifyDown(0);
		}
		return min;
	};
	const heapifyUp = (i) => {
		while (i) {
			const parent = Math.floor((i - 1) / 2);
			if (heap[i][0] < heap[parent][0]) {
				[heap[i], heap[parent]] = [heap[parent], heap[i]];
				i = parent;
			} else {
				break;
			}
		}
	};
	const heapifyDown = (i) => {
		const left = i * 2 + 1;
		const right = i * 2 + 2;
		let min = i;
		if (left < heap.length && heap[left] < heap[i]) {
			min = left;
		}
		if (right < heap.length && heap[right] < heap[i]) {
			min = right;
		}
		if (min !== i) {
			[heap[min], heap[i]] = [heap[i], heap[min]];
			heapifyDown(min);
		}
	};
	// 将第一列元素加入优先队列
	for (let i = 0; i < matrix.length; i++) {
		add([matrix[i][0], i, 0]);
	}

	// 从优先队列中弹出 k - 1 个最小元素
	for (let i = 0; i < k - 1; i++) {
		const [val, m, n] = pop();
		if (n < matrix[0].length - 1) {
			add([matrix[m][n + 1], m, n + 1]);
		}
	}

	// 优先队列的队首元素即为第 k 小的元素
	return heap[0][0];
};

相关题目

题号	标题	题解	标签	难度
373	查找和最小的 K 对数字	[✓]	数组 堆（优先队列）
668	乘法表中第k小的数		数学 二分查找
719	找出第 K 小的数对距离		数组 双指针 二分查找 1+
786	第 K 个最小的质数分数		数组 双指针 二分查找 2+