leetcode215. Kth Largest Element in an Array

时间:2024-10-15 15:51:39

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

Can you solve it without sorting?

Example 1:

Input: nums = [3,2,1,5,6,4], k = 2
Output: 5

Example 2:

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4

Constraints:

1 <= k <= nums.length <= 105
-104 <= nums[i] <= 104

https://leetcode.cn/problems/kth-largest-element-in-an-array/description/
方法一:排序

使用编程语言的内置排序算法对数组 nums 进行排序,然后返回第 N−k 个元素即可。

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return sorted(nums)[len(nums) - k]

方法二:快速选择

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

哨兵划分: 以数组某个元素(一般选取首元素)为基准数,将所有小于基准数的元素移动至其左边,大于基准数的元素移动至其右边。
递归: 对 左子数组 和 右子数组 递归执行 哨兵划分,直至子数组长度为 1 时终止递归,即可完成对整个数组的排序。

下图展示了数组 [2,4,1,0,3,5] 的快速排序流程。

在这里插入图片描述
「快速选择」:设 N 为数组长度。根据快速排序原理,如果某次哨兵划分后,基准数的索引正好是 N−k ,则意味着它就是第 k 大的数字 。此时就可以直接返回它,无需继续递归下去了。

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

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

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

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)