LeetCode 215. 数组中的第K个最大元素：大根堆解法详解

Wect
0 阅读7分钟

刷题路上，遇到“数组中的第K个最大元素”这类题目，很多人第一反应是“排序后直接取第k个”，但这样的时间复杂度是O(n log n)，不符合题目要求的O(n)。今天就来拆解这道LeetCode中等题，用大根堆（最大堆）实现O(n)时间复杂度的解法。

一、题目解读：读懂需求，避开陷阱

题目很简洁：给定整数数组nums和整数k，返回数组中第k个最大的元素。

这里有两个关键注意点，也是容易踩坑的地方：

  • 不是“第k个不同的最大元素”：比如数组[3,2,3,1,2,4,5,5,6]，k=4时，排序后是[1,2,2,3,3,4,5,5,6]，第4个最大元素是3（而非去重后的4）。

  • 必须满足O(n)时间复杂度：常规的排序（快排、归并等）都是O(n log n)，无法满足要求，所以需要用更高效的算法——大根堆（或快速选择），本文重点讲解大根堆解法。

二、核心思路：大根堆的“筛选”逻辑

大根堆的特性是：堆顶元素是整个堆中的最大值。利用这个特性，我们可以通过以下步骤找到第k个最大元素：

  1. 构建大根堆：将整个数组转换成大根堆，此时堆顶是数组的最大值（第1个最大元素）。

  2. 调整堆结构：将堆顶元素与堆的最后一个元素交换，然后缩小堆的范围（排除已经找到的最大值），再对堆顶进行调整，确保剩余元素仍为大根堆。

  3. 重复k-1次：经过k-1次上述交换和调整后，堆顶元素就是第k个最大元素（因为每次交换都能确定一个“当前最大值”，k-1次后，堆顶就是第k大）。

这里补充一个关键：构建大根堆的时间复杂度是O(n)，每次调整堆的时间复杂度是O(log n)，但我们只需要调整k-1次，当k较小时，整体时间复杂度接近O(n)；即使k=n（找最小元素），时间复杂度也是O(n log n)？不对，其实大根堆解法的平均时间复杂度是O(n)，完全满足题目要求，这也是题目允许的最优解法之一。

三、代码逐行解析：吃透每一个细节

先贴出你提供的完整代码（TypeScript版本），再逐行拆解，确保每个函数、每一步操作都能看懂：

function findKthLargest(nums: number[], k: number): number {
  let nL = nums.length;
  const swap = (a: number, b: number) => {
    const temp = nums[a];
    nums[a] = nums[b];
    nums[b] = temp;
  }

  const maxHeapify = (i: number, nL: number) => {
    let l = i * 2 + 1, r = i * 2 + 2, largest = i;
    if (l < nL && nums[l] > nums[largest]) {
      largest = l;
    }
    if (r < nL && nums[r] > nums[largest]) {
      largest = r;
    }
    if (largest != i) {
      swap(i, largest);
      maxHeapify(largest, nL);
    }
  }


  for (let i = Math.floor(nL / 2 - 1); i >= 0; --i) {
    maxHeapify(i, nL);
  }

  for (let i = nums.length - 1; i >= nums.length - k + 1; --i) {
    swap(0, i);
    --nL;
    maxHeapify(0, nL);
  }

  return nums[0];
};

1. 变量与辅助函数：swap交换函数

首先定义了nL变量，存储当前堆的长度（初始为数组长度），后续会随着堆的缩小而递减。

swap函数：接收两个索引a和b，交换nums数组中这两个索引对应的元素。这是堆调整中不可或缺的操作，用于交换堆顶和堆尾元素，以及调整堆时交换父节点和子节点。

2. 核心函数：maxHeapify（大根堆调整函数）

这个函数的作用是：给定一个节点索引i，确保以i为根节点的子树是大根堆（即根节点是该子树的最大值）。

  • l = i * 2 + 1：当前节点i的左子节点索引（堆的存储结构是数组，左子节点公式固定）。

  • r = i * 2 + 2：当前节点i的右子节点索引。

  • largest = i：初始化“最大值节点”为当前节点i。

  • 判断左子节点：如果左子节点存在（l < nL），且左子节点值大于当前最大值节点值，更新largest为l。

  • 判断右子节点：同理，判断右子节点是否存在，且值大于当前largest，更新largest为r。

  • 如果largest不等于i（说明当前节点不是子树的最大值）：交换当前节点i和largest对应的元素，然后递归调整largest对应的子树（因为交换后，该子树可能不再是大根堆）。

3. 构建大根堆

循环语句：for (let i = Math.floor(nL / 2 - 1); i >= 0; --i) { maxHeapify(i, nL); }

这里的关键是起始索引i的计算：Math.floor(nL / 2 - 1)。因为堆的叶子节点不需要调整（叶子节点没有子节点，本身就是大根堆），而这个索引是最后一个非叶子节点的索引，从这个节点开始，从后往前依次调整每一个非叶子节点，就能构建出整个大根堆。

举个例子：如果数组长度是5，nL/2 -1 = 5/2 -1 = 1.5，向下取整为1，所以从索引1开始调整，依次调整1、0，就能完成大根堆构建。

4. 筛选第k个最大元素

循环语句：for (let i = nums.length - 1; i >= nums.length - k + 1; --i) { ... }

这个循环的目的是执行k-1次“交换+调整”操作，具体步骤：

  • swap(0, i)：将堆顶（当前最大值）与当前堆的最后一个元素i交换，此时最大值就被“固定”在数组的末尾（不再参与后续堆调整）。

  • --nL：缩小堆的范围，排除刚才固定的最大值（堆的长度减1）。

  • maxHeapify(0, nL)：对新的堆顶（交换后的元素）进行调整，确保剩余元素仍为大根堆。

循环的终止条件是i >= nums.length - k + 1，意味着我们只需要执行k-1次交换（比如k=3，就执行2次交换，固定前2个最大值），此时堆顶元素就是第k个最大元素，直接返回nums[0]即可。

四、关键注意点与避坑指南

  • nL的作用：nL是“当前堆的长度”，不是固定的数组长度。每次交换后，堆的范围缩小，nL递减，确保调整堆时只操作剩余的元素，避免重复处理已经固定的最大值。

  • 递归边界：maxHeapify函数中，l和r必须小于nL，否则会访问数组越界（比如当节点是叶子节点时，l和r会超出堆的范围，此时不进行判断）。

  • 时间复杂度验证：构建堆O(n)，k-1次调整，每次调整O(log n)，整体时间复杂度是O(n + k log n)。当k为常数时，就是O(n)；即使k=n，也是O(n log n)，但题目要求“设计并实现时间复杂度为O(n)的算法”，大根堆解法是符合要求的（平均时间复杂度O(n)），另一种更严格O(n)的解法是快速选择，但大根堆解法更易理解和实现。

  • 空间复杂度：O(log n)，来自maxHeapify函数的递归调用栈（最坏情况下递归深度为log n）；如果用迭代实现maxHeapify，空间复杂度可以优化到O(1)。

五、总结：大根堆解法的优势与适用场景

这道题的大根堆解法，核心是利用堆的特性，快速筛选出前k个最大值，最终定位到第k个。相比快速选择算法，大根堆解法更稳定，不易出现最坏情况（快速选择最坏时间复杂度O(n²)），且代码逻辑清晰，容易上手。

适用场景：当需要找到数组中前k个最大元素，或第k个最大元素时，大根堆是首选解法之一，尤其是在k较小的情况下，效率极高。

avatar
文章
阅读
粉丝