跳到主要内容

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

题目描述​

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

请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

示例​

示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5

示例 2:

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

解题思路​

第一步:理解问题本质​

第 k 大元素 = 升序排序后下标为 n-k 的元素。

例如:[1,2,3,4,5],第 2 大 = 4,升序下标为 5-2=3。

第二步:暴力解法 - 完全排序​

# 完全排序后按下标取:第 k 大就是升序里倒数第 k 个元素
def findKthLargest(self, nums: List[int], k: int) -> int:
nums.sort()
# 用 n-k 而不是 k-1,因为「第 k 大」是从大往小数
return nums[len(nums) - k]

为什么不够好:时间复杂度 O(n log n),没有利用"只需要第 k 大"这个信息。

第三步:优化解法 - 最小堆​

维护一个大小为 k 的最小堆,堆顶就是第 k 大的元素。

# 只保留「目前最大的 k 个数」,堆顶就是这 k 个里最小的那个,即第 k 大
def findKthLargest(self, nums: List[int], k: int) -> int:
min_heap = []
for num in nums:
# 先入堆,让每个数都有机会参与「前 k 大」的竞争
heapq.heappush(min_heap, num)
# 超过 k 个就把最小的挤出去,堆永远只有 k 个元素
if len(min_heap) > k:
heapq.heappop(min_heap)
# 剩下的这 k 个元素中,最小的那个正是第 k 大
return min_heap[0]

分析:时间 O(n log k),空间 O(k)。

第四步:最优解法 - 快速选择​

快速排序的变种,平均 O(n) 时间找到第 k 大元素。


完整代码实现​

算法可视化 · 逐帧回放这道题的代码91 帧 · 录制输入 [3,2,1,5,6,4], 2
from typing import List
from random import randint

class Solution:
"""
数组中的第K个最大元素 - 快速选择算法

核心思路:
快速选择算法是快速排序的变种,利用 partition 操作在平均 O(n) 时间内找到第 k 大元素。

时间复杂度:平均 O(n),最坏 O(n²)
空间复杂度:O(1)
"""

def partition(self, nums: List[int], left: int, right: int) -> int:
"""划分函数,返回 pivot 的最终位置"""
# 随机选择 pivot
i = randint(left, right)
pivot = nums[i]
nums[i], nums[left] = nums[left], nums[i]

# 双指针划分
i, j = left + 1, right
while True:
while i <= j and nums[i] < pivot:
i += 1
while i <= j and nums[j] > pivot:
j -= 1
if i >= j:
break
nums[i], nums[j] = nums[j], nums[i]
i += 1
j -= 1

nums[left], nums[j] = nums[j], nums[left]
return j

def findKthLargest(self, nums: List[int], k: int) -> int:
n = len(nums)
target_index = n - k # 第 k 大在升序中的下标
left, right = 0, n - 1

while True:
i = self.partition(nums, left, right)
if i == target_index:
return nums[i]
if i > target_index:
right = i - 1
else:
left = i + 1

示例推演​

以 nums = [3,2,1,5,6,4], k = 2 为例:

目标:找第 2 大 = 升序下标 6-2=4 的元素。

第1轮 partition(假设 pivot = 4):

初始: [3,2,1,5,6,4]
交换 pivot 到首位: [4,2,1,5,6,3]
划分后: [3,2,1,4,6,5]
pivot 位置: 3

3 < 4(target_index),在右半部分继续。

第2轮 partition(在 [6,5] 中,假设 pivot = 6):

[6,5] -> 划分后 [5,6]
pivot 位置: 5(相对于原数组是 5)

5 > 4(target_index),目标在左半部分继续。

第3轮 partition(在 [5] 中):

pivot = 5,位置: 4

4 == 4,找到!返回 nums[4] = 5。


复杂度分析​

解法时间复杂度空间复杂度说明
完全排序O(n log n)O(1) 或 O(n)取决于排序算法
最小堆O(n log k)O(k)适合 k 较小的情况
快速选择平均 O(n),最坏 O(n²)O(1)最优解法

快速选择平均 O(n) 的证明:

T(n) = n + n/2 + n/4 + ... = 2n = O(n)


易错点总结​

1. target_index 计算​

错误:target_index = k - 1

正确:target_index = n - k(第 k 大 = 升序下标 n-k)

2. partition 后的判断​

if i == target_index:
# pivot 已经落到最终位置,正好是目标下标,直接返回
return nums[i]
if i > target_index: # 注意是 > 不是 <
# pivot 在目标右侧,答案只可能在左半区间
right = i - 1
else:
# 反之答案在右半区间,每轮丢掉一半,所以平均是 O(n)
left = i + 1

3. 随机 pivot 的重要性​

如果不随机选择 pivot,在已排序数组上会退化到 O(n²)。


扩展思考​

1. 找中位数​

k = n/2,本题特例。

2. 找前 k 大元素​

347. 前 K 个高频元素,类似思路。

3. 相关题目​


相关题目​

自测​

加载评论中...