快速选择算法

1 问题描述

适用于求数组中的第K个最大元素
示例： 215. 数组中的第K个最大元素

2 解题思路

2.1 快速排序

2.1.1 算法思路

快速排序的核心包括“哨兵划分” 和 “递归” 。

• 哨兵划分： 以数组某个元素（一般选取首元素）为基准数，将所有小于基准数的元素移动至其左边，大于基准数的元素移动至其右边。
• 递归： 对 左子数组 和 右子数组 递归执行 哨兵划分，直至子数组长度为 1 时终止递归，即可完成对整个数组的排序。
  下图展示了数组[2,4,1,0,3,5]的快速排序流程。
  
  「快速选择」：设 N为数组长度。根据快速排序原理，如果某次哨兵划分后，基准数的索引正好是 $N?k$ ，则意味着它就是第 kkk 大的数字 。此时就可以直接返回它，无需继续递归下去了。

然而，对于包含大量重复元素的数组，每轮的哨兵划分都可能将数组划分为长度为 1 和 n?1 的两个部分，这种情况下快速排序的时间复杂度会退化至 $O(N^2)$
) 。

一种解决方案是使用「三路划分」，即每轮将数组划分为三个部分：小于、等于和大于基准数的所有元素。这样当发现第 k大数字处在“等于基准数”的子数组中时，便可以直接返回该元素。

为了进一步提升算法的稳健性，我们采用随机选择的方式来选定基准数。

2.1.2 代码示例

class Solution:
    def findKthLargest(self, nums, k):
        def quick_select(nums, k):
            # 随机选择基准数
            pivot = random.choice(nums)
            big, equal, small = [], [], []
            # 将大于、小于、等于 pivot 的元素划分至 big, small, equal 中
            for num in nums:
                if num > pivot:
                    big.append(num)
                elif num < pivot:
                    small.append(num)
                else:
                    equal.append(num)
            if k <= len(big):
                # 第 k 大元素在 big 中，递归划分
                return quick_select(big, k)
            if len(nums) - len(small) < k:
                # 第 k 大元素在 small 中，递归划分
                return quick_select(small, k - len(nums) + len(small))
            # 第 k 大元素在 equal 中，直接返回 pivot
            return pivot
        
        return quick_select(nums, k)

2.2 基于堆排序的选择方法

2.2.1 算法思路

• 可以先构造一个空间大小为k的堆
• 再从nums[k+1]开始与堆顶元素比较
• 如果nums[i] > heap则将当前堆顶元素出堆,并将nums[i]入堆
• 这样遍历完数组后堆顶元素就是ans

2.2.2 代码示例

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # 先将前k个元素入堆
        heap = [nums[i] for i in range(0,k)]
        heapq.heapify(heap)
        for i in range(k,len(nums)):
            if nums[i] > heap[0]:
                heapq.heappop(heap) # 弹出堆顶元素
                heapq.heappush(heap,nums[i]) # nums[i] 入堆
        return heap[0]

参考资料：
https://leetcode.cn/problems/kth-largest-element-in-an-array/solutions/2361969/215-shu-zu-zhong-de-di-k-ge-zui-da-yuan-d786p/
https://leetcode.cn/problems/kth-largest-element-in-an-array/solutions/2559996/li-yong-dui-jie-jue-top-kwen-ti-by-kiski-1747/