215. 数组中的第 K 个最大元素 (Medium)
专题归类: 06-栈与堆 · 09-二分查找 LeetCode 链接: https://leetcode.cn/problems/kth-largest-element-in-an-array/
在线做题: 写 Python 代码并运行 · Java 跳转 LeetCode
题目描述
给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。
请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。
你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。
示例 1:
输入:nums = [3,2,1,5,6,4], k = 2
输出:5
示例 2:
输入:nums = [3,2,3,1,2,4,5,5,6], k = 4
输出:4
题目详细分析
- 数据范围: 1 ≤ k ≤ nums.length ≤ 10^5,-10^4 ≤ nums[i] ≤ 10^4。数组最大 10^5 长度,O(n log n) 可以接受但题目要求 O(n)。
- 输入输出特征: 找第 k 大(从大到小排序后第 k 个),不是第 k 个不同的元素。可能有重复值。
- 边界条件:
- k = 1 → 返回最大值。
- k = n → 返回最小值。
- 所有元素相同 → 返回该值。
- 核心约束: 时间复杂度 O(n)。这排除了直接排序(O(n log n))的方案。
- 隐藏条件:
- 第 k 大 = 第 n-k+1 小(从小到大排序后的索引)。
- 不需要对整个数组排序,只需要找到第 k 大的元素。
- 快速选择(Quick Select)是满足 O(n) 要求的经典算法。
- 最小堆(大小为 k)也是常见解法,时间复杂度 O(n log k),对于 k << n 时非常实用。
小白版直白理解
就像在一堆考卷中找出 第 k 高的分数。正常的做法是全部排序然后取第 k 个(O(n log n)),但你实际上不需要全部排好。
两种”偷懒”方法:
方法一(最小堆): 准备一个只能装 k 个人的领奖台。每次来一个人,如果台上的最矮的人比新人矮,就把最矮的踢下去,新人上台。遍历完所有人后,领奖台最矮的那个就是第 k 高的(因为台上是前 k 高的人)。
方法二(快速选择): 就像在人群中寻找第 k 高的人——随便指一个人作为标杆,让所有比他高的人站左边,比他矮的站右边。看标杆现在是第几高:如果正好是第 k,就找到了;如果排名太靠前(更高),就只在左边找;如果排名太靠后(更矮),就只在右边找。每次能排除一半人,所以平均 O(n)。
解题思路
思路一:最小堆(适合 k 较小的情况)(推荐)
核心思想: 维护一个大小为 k 的最小堆。遍历数组元素,将元素加入堆中;当堆大小 > k 时,弹出堆顶(当前最小的元素)。遍历结束后,堆顶就是第 k 大的元素。
为什么最小堆可以找到第 k 大?堆中始终保存着当前最大的 k 个元素,堆顶是这 k 个中最小的一即第 k 大的。
import heapq
def findKthLargest(nums, k):
heap = []
for num in nums:
heapq.heappush(heap, num) # 入堆
if len(heap) > k:
heapq.heappop(heap) # 保持堆大小为 k
return heap[0] # 堆顶就是第 k 大复杂度: 时间 O(n log k),空间 O(k)。当 k << n 时非常高效。
思路二:快速选择(Quick Select)—— 平均 O(n)
核心思想: 基于快速排序的 partition 操作。随机选择一个 pivot,将数组分为大于 pivot 和小于 pivot 两部分。如果 pivot 恰好是第 k-1 个(0-index 从大到小),返回 pivot;否则只在包含目标的一侧递归。
import random
def findKthLargest(nums, k):
def partition(left, right):
"""将区间 [left, right] 按 pivot 分区,返回 pivot 最终位置"""
pivot_idx = random.randint(left, right) # 随机选 pivot
nums[pivot_idx], nums[right] = nums[right], nums[pivot_idx] # 放到最右边
pivot = nums[right]
i = left
for j in range(left, right):
if nums[j] > pivot: # 大于 pivot 的放左边
nums[i], nums[j] = nums[j], nums[i]
i += 1
nums[i], nums[right] = nums[right], nums[i]
return i
left, right = 0, len(nums) - 1
k -= 1 # 转换为 0-index
while True:
pos = partition(left, right)
if pos == k:
return nums[pos]
elif pos < k:
left = pos + 1 # 在右半部分找
else:
right = pos - 1 # 在左半部分找复杂度: 平均 O(n),最坏 O(n²)(当每次 pivot 都选到最值时)。时间复杂度分析:n + n/2 + n/4 + … ≈ 2n = O(n)。
思路三:直接排序(不满足 O(n) 要求,但最简洁)
def findKthLargest(nums, k):
return sorted(nums, reverse=True)[k - 1]复杂度: O(n log n),不满足题目 O(n) 要求但实际面试中有时可先提出来作为 baseline。
易错点
- 第 k 大 vs 第 k 小: 第 k 大是从大到小排序后的第 k 个。如果使用从小到大排序,要取
nums[n - k]。快速选择中如果 partition 按大于 pivot 的放左边,则 pos 表示从大到小的排名(0-index)。 - k 转换为 0-index: 如果 partition 返回的是从大到小排序的 0-index 位置,需要将 k 减 1 再比较。
- partition 方向一致性: 分区时如果大于 pivot 的放左边,那么 partition 返回的 pos 就是该元素在从大到小排序中的位置。确保比较逻辑与此一致。
- 重复元素: 有重复元素时 partition 可能不如预期稳定。例如
[3,2,3,1,2,4,5,5,6],k=4,partition 时相等的元素可以放在左侧或右侧,不影响最终结果。 - 随机化的重要性: 如果不随机选择 pivot,最坏情况(已排序数组)下快速选择会退化到 O(n²)。随机 pivot 能保证期望 O(n)。
框架提炼
快速选择模板:
import random
def quick_select(nums, k):
"""返回数组中第 k 大元素(k 从 1 开始)"""
def partition(l, r):
pivot = nums[random.randint(l, r)]
# 将 pivot 放到最右边
p_idx = l
for i in range(l, r + 1):
if nums[i] == pivot:
p_idx = i
break
nums[p_idx], nums[r] = nums[r], nums[p_idx]
i = l
for j in range(l, r):
if nums[j] > pivot: # 大于 pivot 放左边
nums[i], nums[j] = nums[j], nums[i]
i += 1
nums[i], nums[r] = nums[r], nums[i]
return i
l, r = 0, len(nums) - 1
k -= 1 # 0-index
while True:
pos = partition(l, r)
if pos == k:
return nums[pos]
elif pos < k:
l = pos + 1
else:
r = pos - 1Top K 问题常用方案对比:
| 方法 | 时间复杂度 | 空间复杂度 | 适用场景 |
|---|---|---|---|
| 快速选择 | O(n) 平均 | O(1) | 一次查找,需要 O(n) |
| 最小堆 | O(n log k) | O(k) | k 较小,或数据流场景 |
| 排序 | O(n log n) | O(1) | 需要全部有序时 |
| 桶排序 | O(n) | O(n) | 数据范围较小时 |
关联题目
- 347-前K个高频元素 — 堆的应用进阶,多了一层频次统计。同样是 Top K 问题,但需要先统计频率再用堆选 Top K。
- 295-数据流的中位数 — 动态数据流中的第 k 大元素(k = n/2),使用双堆技巧。与 215 的静态数组不同,数据流要求动态维护。
- 703-数据流中的第K大元素 — 215 的在线版本:不断有新的元素加入,需要随时知道第 k 大的值。使用最小堆维护大小为 k 的窗口。