【LeetCode 热题 100】215. 数组中的第K个最大元素（Python 快速选择详解）

在刷 LeetCode 的过程中，“第K大”是一个非常高频的考点，而题目 215. 数组中的第K个最大元素 就是经典代表。这道题不仅考察我们对排序的理解，还挑战我们写出时间复杂度为 O(n) 的算法。

本文将带你深入理解并实现一个基于快速选择（Quickselect）的高性能解法。

🧩 题目描述

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

⚠️ 注意：题目中要求是第 k 大，而不是第 k 个不同的元素。

示例

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

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

最简单的方法就是对数组排序后取倒数第 k 个元素：

def findKthLargest(nums, k):
    nums.sort()
    return nums[-k]

虽然代码简洁，但排序的时间复杂度是 O(n log n)，不符合题目期望的 O(n) 要求。

💡 正确姿势：快速选择算法（Quickselect）

📚 Quickselect 是什么？

快速选择是快速排序的“变种”，它利用了分治的思想，在每次划分时只处理可能包含答案的那一半，从而平均时间复杂度降为 O(n)。

🧠 算法思路

1. 随机选一个 pivot（基准元素）
2. 将数组划分为三部分：
   • big：所有大于 pivot 的数
   • equal：所有等于 pivot 的数
   • small：所有小于 pivot 的数
3. 判断第 k 大数在哪个部分：
   • 如果在 big 中，递归查找 big 中的第 k 大
   • 如果在 equal 中，直接返回 pivot
   • 如果在 small 中，调整 k，递归查找 small 中的 (k - big数 - equal数) 大

✅ 完整 Python 实现

import random

class Solution:
    def findKthLargest(self, nums, k):
        def quick_select(nums, k):
            pivot = random.choice(nums)
            big = [num for num in nums if num > pivot]
            small = [num for num in nums if num < pivot]
            equal_count = len(nums) - len(big) - len(small)

            if k <= len(big):
                return quick_select(big, k)
            elif k <= len(big) + equal_count:
                return pivot
            else:
                return quick_select(small, k - len(big) - equal_count)

        return quick_select(nums, k)

🔍 递归是怎么用的？

• 递归本质：自己调用自己，逐步缩小问题范围。
• 每次我们都在更小的数组中递归地寻找第 k 大元素。
• 有效缩小搜索空间，每次最多处理一半元素。

🎯 举个例子

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

• 首轮 pivot 可能是 3：
  • big = [5, 6, 4]，长度为 3
• 因为 2 ≤ 3，说明我们要找的第 2 大在 big 中
• 递归进入 quick_select([5, 6, 4], 2)
• 再选 pivot，比如是 5：
  • big = [6]，equal = [5]，small = [4]
• 此时 2 落在 equal 区间，返回 5，找到答案！

⚙️ 时间 & 空间复杂度分析

🧠 优化建议

1. 节省空间：可以用原地 partition（如 Lomuto 法）替代切片，避免生成新列表。
2. 面试场景：如果面试官没要求 O(n)，直接用排序即可快速出解。
3. 重复元素处理：等于 pivot 的元素不能漏掉，需要单独统计。

📌 总结

• 快速选择是解决“第K大”、“第K小”类问题的利器
• 随机选择 pivot 是算法性能稳定的关键
• 递归思想+分治策略，让问题逐步缩小，最终得到答案