算法学习

Appearance

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

给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。

请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。

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

示例 1:
输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
解释：矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13

示例 2:
输入：matrix = [[-5]], k = 1
输出：-5

解题思路 ​

输入： 给你一个 n x n 矩阵 matrix

输出： 找到矩阵中第 k 小的元素

方法分析 ​

  • 我们可以使用 最小堆（Min Heap） 来降低复杂度。

具体做法 ​

  1. 初始化堆：
    • 把每一行的第一个元素 (值, 行号, 列号) 入堆。
    • 例如第 0 行的第一个元素是 (matrix[0][0], 0, 0)。
    • 堆中存储结构为 [value, row, col]。
  2. 重复弹出 k 次：
    • 每次从堆中弹出最小值 (val, r, c)；
    • 如果这一行还有下一个元素（即 c + 1 < n），就把 (matrix[r][c + 1], r, c + 1) 入堆；
    • 重复该过程 k 次，第 k 次弹出的元素就是答案。

复杂度分析 ​

  • 堆最多存 n 个元素，每次弹出/插入是 log n
  • 总时间复杂度：O(k log n)
  • 空间复杂度：O(n)

代码实现 ​

python
import heapq

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        """
        思路：最小堆（Min Heap）——多路归并法
        每一行都是升序的，把每一行看作一个有序数组，
        然后使用最小堆不断取出当前最小值。
        """

        n = len(matrix)
        heap = []  # 定义最小堆，元素格式为 (值, 行号, 列号)

        # 1️⃣ 初始化：将每行的第一个元素入堆
        for i in range(n):
            # 将 (matrix[i][0], i, 0) 放入堆中
            # 含义：值 = matrix[i][0]，来自第 i 行，第 0 列
            heapq.heappush(heap, (matrix[i][0], i, 0))

        # 2️⃣ 弹出最小值 k-1 次，第 k 个最小值就是答案
        for _ in range(k - 1):
            val, r, c = heapq.heappop(heap)  # 弹出当前堆顶（最小值）

            # 如果该元素所在行还有下一个元素，则入堆
            if c + 1 < n:
                heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
        
        # 3️⃣ 此时堆顶就是第 k 小的元素
        return heap[0][0]
javascript
class MyMinHeap {
  constructor() {
    this.data = [];
  }
  size() { return this.data.length; }
  peek() { return this.data[0]; }
  push(val) {
    this.data.push(val);
    this._heapifyUp();
  }
  pop() {
    if (this.size() === 0) return null;
    const top = this.data[0];
    const end = this.data.pop();
    if (this.size() > 0) {
      this.data[0] = end;
      this._heapifyDown();
    }
    return top;
  }
  _heapifyUp() {
    let i = this.size() - 1;
    while (i > 0) {
      const p = Math.floor((i - 1) / 2);
      if (this.data[p][0] <= this.data[i][0]) break;
      [this.data[p], this.data[i]] = [this.data[i], this.data[p]];
      i = p;
    }
  }
  _heapifyDown() {
    let i = 0;
    const n = this.size();
    while (true) {
      const l = 2 * i + 1, r = 2 * i + 2;
      let smallest = i;
      if (l < n && this.data[l][0] < this.data[smallest][0]) smallest = l;
      if (r < n && this.data[r][0] < this.data[smallest][0]) smallest = r;
      if (smallest === i) break;
      [this.data[i], this.data[smallest]] = [this.data[smallest], this.data[i]];
      i = smallest;
    }
  }
}

var kthSmallest = function(matrix, k) {
  const n = matrix.length;
  const heap = new MyMinHeap();

  // 初始化：每行第一个元素入堆
  for (let i = 0; i < n; i++) {
    heap.push([matrix[i][0], i, 0]);
  }

  let result = null;
  // 弹出 k 次
  for (let i = 0; i < k; i++) {
    const [val, r, c] = heap.pop();
    result = val;
    if (c + 1 < n) heap.push([matrix[r][c + 1], r, c + 1]);
  }

  return result;
};

复杂度分析 ​

时间复杂度：O(k log n)

空间复杂度：O(n)

链接 ​

378 国际版

378 中文版