215. 数组中的第 K 个最大元素

LeetCode 原题链接

1. 题目描述

给定一个整数数组 nums 和一个整数 k,返回数组中第 k 个最大的元素。

注意这里要找的是排序后的第 k 大元素,不是第 k 个不同的元素。也就是说,重复元素会参与排名。

示例 1:

输入:nums = [3, 2, 1, 5, 6, 4], k = 2
输出:5
解释:降序排列后是 [6, 5, 4, 3, 2, 1],第 2 大元素是 5。

示例 2:

输入:nums = [3, 2, 3, 1, 2, 4, 5, 5, 6], k = 4
输出:4
解释:降序排列后是 [6, 5, 5, 4, 3, 3, 2, 2, 1],第 4 大元素是 4。

2. 题型判断

本题要求在无序数组中找到第 k 个最大的元素,属于典型的 Top K 问题。

直观做法是先排序,再取下标 k - 1 的元素,时间复杂度是 O(n log n)。但题目只需要第 k 大,不需要整个数组有序,因此可以用快速选择。快速选择每次通过一次分区确定一个元素的最终排名,再只递归处理可能包含答案的一侧,平均时间复杂度可以降到 O(n)

3. 核心思路

使用三路快速选择。每次随机选择一个 pivot,把当前搜索区间原地划分为三段:

  • 左侧:大于 pivot
  • 中间:等于 pivot
  • 右侧:小于 pivot

设中间等值区间为 [lt, gt]

  • 如果 targetIndex < lt,答案在左侧;
  • 如果 targetIndex > gt,答案在右侧;
  • 如果 lt <= targetIndex <= gt,目标排名落在等值区间内,直接返回 pivot

与两路分区相比,三路分区可以一次跳过全部重复的 pivot。对于 [2,2,2,2,2] 这类数组,一轮分区就能结束,不会每轮只排除一个元素而退化超时。

4. 步骤图解

下面以 nums = [3, 2, 3, 1, 2, 4, 5, 5, 6]k = 4 为例。目标下标为 targetIndex = 3

假设本轮随机选到 pivot = 5,三路分区后数组在逻辑上变成:

大于 5        等于 5           小于 5
[6]        |  [5, 5]  |  [其余元素]
下标 0        下标 1~2          下标 3~8

等值区间为 [1, 2],目标下标 3 > 2,所以只需继续搜索右侧 [3, 8]

假设下一轮选到 pivot = 4

大于 4             等于 4       小于 4
[6, 5, 5]       |    [4]    |  [3, 3, 2, 2, 1]
下标 0~2            下标 3       下标 4~8

目标下标 3 落在等值区间 [3, 3] 中,所以答案为 4

5. 变量与区间含义

  • targetIndex:第 k 大元素在降序数组中的下标,也就是 k - 1
  • left:当前查找区间的左边界。
  • right:当前查找区间的右边界。
  • pivot:当前区间随机选择的基准值。
  • lt:等值区间的左边界;[left, lt - 1] 都大于 pivot
  • i:当前正在检查的元素下标。
  • gt:等值区间的右边界;[gt + 1, right] 都小于 pivot

快速选择维护闭区间 [left, right]。答案一定在该区间内,每次分区后根据等值区间 [lt, gt]targetIndex 的关系缩小范围。

6. 推进规则

分区过程中始终维护:

[left, lt - 1]  > pivot
[lt, i - 1]     = pivot
[i, gt]         尚未处理
[gt + 1, right] < pivot

根据 nums[i] 分三种情况:

  1. nums[i] > pivot:与 nums[lt] 交换,lt++i++
  2. nums[i] < pivot:与 nums[gt] 交换,gt--,但 i 不动,因为换过来的元素尚未检查;
  3. nums[i] === pivot:元素留在中间,i++

查找规则:

  • 如果 targetIndex < lt,更新 right = lt - 1
  • 如果 targetIndex > gt,更新 left = gt + 1
  • 否则目标下标位于等值区间,直接返回 pivot

7. 边界条件与易错点

  • k 是从 1 开始计数的,数组下标从 0 开始,所以目标下标是 k - 1
  • 分区按降序处理,所以大于 pivot 的元素放左边,小于 pivot 的元素放右边。
  • nums[i] < pivot 时,交换到 i 的元素还没有分类,因此不能立即执行 i++
  • 必须返回整个等值区间 [lt, gt],不能只返回某个 pivot 下标,否则会失去三路分区处理重复值的优势。
  • 快速选择会修改数组顺序。如果不希望修改原数组,可以先复制一份数组再处理。
  • 随机选择 pivot 可以降低连续选到极端值、退化成 O(n²) 的概率。

8. 代码实现

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

  while (left <= right) {
    // 三路分区返回所有等于 pivot 的元素所在区间 [lt, gt]。
    const [lt, gt] = threeWayPartition(nums, left, right);

    if (targetIndex < lt) {
      right = lt - 1;
    } else if (targetIndex > gt) {
      left = gt + 1;
    } else {
      // 目标排名落在等值区间内,该区间的任意元素都是答案。
      return nums[targetIndex];
    }
  }
};

function threeWayPartition(nums, left, right) {
  // 在闭区间 [left, right] 中随机选择 pivot。
  const pivotIndex = left + Math.floor(Math.random() * (right - left + 1));
  const pivot = nums[pivotIndex];

  // [left, lt - 1]:大于 pivot
  // [lt, i - 1]:等于 pivot
  // [i, gt]:尚未处理
  // [gt + 1, right]:小于 pivot
  let lt = left;
  let i = left;
  let gt = right;

  while (i <= gt) {
    if (nums[i] > pivot) {
      [nums[lt], nums[i]] = [nums[i], nums[lt]];
      lt++;
      i++;
    } else if (nums[i] < pivot) {
      // 换到 i 的元素尚未检查,所以此处不能执行 i++。
      [nums[i], nums[gt]] = [nums[gt], nums[i]];
      gt--;
    } else {
      i++;
    }
  }

  return [lt, gt];
}

9. 其他解法

解法一:排序

最直观的做法是把数组按降序排序,然后返回下标 k - 1 的元素。

这种写法最短,也很适合先解释题意,但它会把整个数组都排好序,而题目只需要第 k 大元素,所以时间复杂度比快速选择更高。

var findKthLargest = function (nums, k) {
  // 降序排序后,第 k 大元素就在下标 k - 1。
  nums.sort((a, b) => b - a);
  return nums[k - 1];
};

解法二:手写快排

也可以手写快速排序,把数组整体按降序排好,再返回 nums[k - 1]

这种解法和内置排序的思想一样,都会排序整个数组。它的优势是能练习快排分区逻辑,但对本题来说不如快速选择高效,因为快速选择只处理可能包含第 k 大的一侧,快排则需要继续排序两侧。

var findKthLargest = function (nums, k) {
  quickSort(nums, 0, nums.length - 1);
  return nums[k - 1];
};

function quickSort(nums, left, right) {
  if (left >= right) {
    return;
  }

  const pivotIndex = partition(nums, left, right);

  quickSort(nums, left, pivotIndex - 1);
  quickSort(nums, pivotIndex + 1, right);
}

function partition(nums, left, right) {
  // 随机选择 pivot,减少有序数组导致退化的概率。
  const pivotIndex = left + Math.floor(Math.random() * (right - left + 1));
  [nums[pivotIndex], nums[right]] = [nums[right], nums[pivotIndex]];

  const pivot = nums[right];
  let storeIndex = left;

  for (let i = left; i < right; i++) {
    // 降序排序:比 pivot 大的元素放左边。
    if (nums[i] > pivot) {
      [nums[storeIndex], nums[i]] = [nums[i], nums[storeIndex]];
      storeIndex++;
    }
  }

  [nums[storeIndex], nums[right]] = [nums[right], nums[storeIndex]];
  return storeIndex;
}

解法三:两路快速选择

两路快速选择使用两个指针从区间两端向中间扫描:

  • i 从左向右寻找小于等于 pivot 的元素;
  • j 从右向左寻找大于等于 pivot 的元素;
  • i <= j 时交换两个元素,使较大的元素靠左、较小的元素靠右。

扫描结束后,把 pivotnums[j] 交换。此时 pivot 位于最终下标 j,左侧元素都大于等于它,右侧元素都小于等于它。然后比较 j 与目标下标,只继续处理可能包含答案的一侧。

与单向分区相比,两路分区遇到等于 pivot 的元素时会停止扫描并交换,再同时移动两个指针,因此可以把重复元素较均匀地分散到两侧,避免全相等数组每轮只排除一个元素。

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

  while (left <= right) {
    const pivotIndex = twoWayPartition(nums, left, right);

    if (pivotIndex === targetIndex) {
      return nums[pivotIndex];
    }

    if (pivotIndex > targetIndex) {
      right = pivotIndex - 1;
    } else {
      left = pivotIndex + 1;
    }
  }
};

function twoWayPartition(nums, left, right) {
  // 随机选择 pivot,再把它暂存到区间左端。
  const randomIndex = left + Math.floor(Math.random() * (right - left + 1));
  [nums[left], nums[randomIndex]] = [nums[randomIndex], nums[left]];

  const pivot = nums[left];
  let i = left + 1;
  let j = right;

  while (true) {
    // 降序分区:左侧保留大于 pivot 的元素。
    while (i <= right && nums[i] > pivot) {
      i++;
    }

    // 右侧保留小于 pivot 的元素。
    while (j >= left + 1 && nums[j] < pivot) {
      j--;
    }

    if (i > j) {
      break;
    }

    [nums[i], nums[j]] = [nums[j], nums[i]];
    i++;
    j--;
  }

  // j 是 pivot 的最终位置。
  [nums[left], nums[j]] = [nums[j], nums[left]];
  return j;
}

例如区间为 [3, 2, 3, 1, 2, 4],选取 pivot = 3。两路扫描结束后,分区结果可能是 [4, 3, 3, 1, 2, 2]pivot 左侧都大于等于 3,右侧都小于等于 3。注意分区只保证两侧与 pivot 的大小关系,并不保证两侧内部有序。

解法四:递归快速选择

快速选择也可以写成递归形式。每次选出一个 pivot,使用三路分区把数组分成三部分:大于 pivot、等于 pivot 和小于 pivot

由于本题寻找的是第 k 大元素,所以优先判断目标排名所在的区域:

  • 如果 k 不超过“大于区域”的长度,就递归查找大于区域。
  • 如果 k 落在“大于区域 + 等于区域”内,答案就是 pivot
  • 否则递归查找小于区域,并扣除已经跳过的元素数量。
/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var findKthLargest = function (nums, k) {
  return quickSelect(nums, k);
};

function quickSelect(nums, k) {
  // 递归到只剩一个元素时,它就是当前区间的第 1 大元素。
  if (nums.length === 1) {
    return nums[0];
  }

  const pivot = nums[Math.floor(Math.random() * nums.length)];
  const greater = [];
  const equal = [];
  const less = [];

  for (const num of nums) {
    if (num > pivot) {
      greater.push(num);
    } else if (num < pivot) {
      less.push(num);
    } else {
      equal.push(num);
    }
  }

  if (k <= greater.length) {
    return quickSelect(greater, k);
  }

  if (k <= greater.length + equal.length) {
    return pivot;
  }

  // 跳过大于 pivot 和等于 pivot 的元素后,更新 k 的相对排名。
  return quickSelect(
    less,
    k - greater.length - equal.length
  );
}

这种写法更直观地体现了递归定义:在当前数组中找到第 k 大元素,就是在分区后可能包含答案的子数组中继续寻找对应排名的元素。为了让分区逻辑更容易理解,代码创建了三个新数组,因此空间开销高于原地快速选择。

解法五:最小堆

维护一个大小为 k 的最小堆。堆里始终保存当前见过的前 k 大元素,堆顶就是这 k 个元素里最小的那个,也就是当前的第 k 大候选值。

遍历数组时:

  • 如果堆的大小小于 k,直接入堆。
  • 如果当前元素大于堆顶,说明它应该进入前 k 大,弹出堆顶后再入堆。
  • 如果当前元素小于等于堆顶,说明它进不了前 k 大,跳过。

遍历结束后,堆顶就是第 k 大元素。

var findKthLargest = function (nums, k) {
  const heap = new MinHeap();

  for (const num of nums) {
    if (heap.size() < k) {
      heap.push(num);
    } else if (num > heap.peek()) {
      heap.pop();
      heap.push(num);
    }
  }

  return heap.peek();
};

class MinHeap {
  constructor() {
    this.heap = [];
  }

  size() {
    return this.heap.length;
  }

  peek() {
    return this.heap[0];
  }

  push(value) {
    this.heap.push(value);
    this.shiftUp(this.heap.length - 1);
  }

  pop() {
    const top = this.heap[0];
    const last = this.heap.pop();

    if (this.heap.length > 0) {
      this.heap[0] = last;
      this.shiftDown(0);
    }

    return top;
  }

  shiftUp(index) {
    while (index > 0) {
      const parent = Math.floor((index - 1) / 2);

      if (this.heap[parent] <= this.heap[index]) {
        break;
      }

      [this.heap[parent], this.heap[index]] = [this.heap[index], this.heap[parent]];
      index = parent;
    }
  }

  shiftDown(index) {
    const n = this.heap.length;

    while (true) {
      let smallest = index;
      const left = index * 2 + 1;
      const right = index * 2 + 2;

      if (left < n && this.heap[left] < this.heap[smallest]) {
        smallest = left;
      }

      if (right < n && this.heap[right] < this.heap[smallest]) {
        smallest = right;
      }

      if (smallest === index) {
        break;
      }

      [this.heap[index], this.heap[smallest]] = [this.heap[smallest], this.heap[index]];
      index = smallest;
    }
  }
}

最小堆适合数据流场景:如果数字一个个到来,不能一次性拿到完整数组,就可以持续维护大小为 k 的堆。

解法六:计数法

如果题目给出的数值范围很小,也可以统计每个数字出现次数,然后从大到小累计数量,累计到 k 时返回当前数字。

这种方法在值域很小时很快,但如果数值范围很大,例如从 -10^910^9,就不适合直接开数组计数。

var findKthLargest = function (nums, k) {
  const countMap = new Map();

  for (const num of nums) {
    countMap.set(num, (countMap.get(num) || 0) + 1);
  }

  const values = [...countMap.keys()].sort((a, b) => b - a);

  for (const value of values) {
    k -= countMap.get(value);

    if (k <= 0) {
      return value;
    }
  }
};

这里用 Map 后仍然需要对不同的值排序,所以复杂度取决于不同数字的个数。如果值域小到可以直接开计数数组,就可以进一步减少排序成本。

10. 复杂度分析

三路快速选择(主解法):

  • 时间复杂度:平均 O(n)。每次分区只处理当前区间,并且平均只继续查找一半区间。
  • 最坏时间复杂度:O(n^2)。如果每次随机选到的 pivot 都接近当前区间的极值,仍可能退化成单边查找;随机选择会显著降低这种情况的概率。
  • 空间复杂度:O(1)。迭代写法只使用常数个额外变量。
  • 重复元素很多时,等于 pivot 的元素会被一次排除,不会像两路分区一样每轮只缩小一个位置。

排序写法:

  • 时间复杂度:O(n log n)
  • 空间复杂度:取决于 JavaScript 引擎的排序实现。

手写快排:

  • 平均时间复杂度:O(n log n)
  • 最坏时间复杂度:O(n^2)。随机选择 pivot 可以降低退化概率。
  • 空间复杂度:平均 O(log n),主要来自递归调用栈;最坏情况下是 O(n)

两路快速选择:

  • 平均时间复杂度:O(n)。每轮分区为线性时间,平均只继续处理一侧。
  • 最坏时间复杂度:O(n^2)。随机选择 pivot 可以降低连续产生极不平衡分区的概率。
  • 空间复杂度:O(1)。代码使用迭代查找和原地分区。
  • 相同元素较多时,两路扫描会让等于 pivot 的元素分散到两侧;但它不能像三路分区一样一次跳过整个等值区间。

递归快速选择:

  • 平均时间复杂度:O(n)。每层只遍历当前数组一次,平均情况下递归处理的数组规模会明显缩小。
  • 最坏时间复杂度:O(n^2)。如果每次选到的 pivot 都接近极值,递归区间会只缩小一个元素。
  • 空间复杂度:平均 O(n),主要来自每层创建的 greaterequalless 数组;递归调用栈平均为 O(log n),最坏为 O(n)

最小堆:

  • 时间复杂度:O(n log k)。堆的大小最多是 k,每次入堆或出堆需要 O(log k)
  • 空间复杂度:O(k)

计数法:

  • 使用 Map 加排序时,时间复杂度是 O(n + m log m),其中 m 是不同数字的个数。
  • 空间复杂度是 O(m)