leetcode-215. 数组中的第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

解题思路

快速选择法
仿照快排，比pivot大的数字放一堆，比pivot小的数字放一堆，相同的数字放一堆，然后分析：

1. larger的数字>=k个，说明需要在larger中继续找
2. larger+equal的数字少于k个，说明需要在smaller中找，因为已经去掉了len(larger) + len(equals)个数字，所以找k - len(larger) + len(equals)个
3. 否则，结果就是pivot

时间复杂度$o(n\log n)$，最好每次选择pivot的时候都用random选，不然无法达到最优时间复杂度
Space complexity: $o(n)$

堆法
保存当前数组中最大的k个数，然后返回最小的那个即可
时间复杂度$o(klogk)$，空间复杂度$o(k)$

代码

快速选择法

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        pivot = random.choice(nums)
        smallers = [item for item in nums if item < pivot]
        equals = [item for item in nums if item == pivot]
        biggers = [item for item in nums if item > pivot]
        if len(biggers) >= k:
            return self.findKthLargest(biggers, k)
        elif len(biggers) + len(equals) < k:
            return self.findKthLargest(smallers, k - len(biggers) - len(equals))
        else:
            return pivot

Without extra space:

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        def helper(left: int, right: int, k: int) -> int:
            if left == right:
                return nums[left]
            pivot = nums[random.randint(left, right)]
            small_index, large_index = left, right
            p = small_index
            while p <= large_index:
                if nums[p] < pivot:
                    nums[small_index], nums[p] = nums[p], nums[small_index]
                    small_index += 1
                    p += 1
                elif nums[p] > pivot:
                    nums[large_index], nums[p] = nums[p], nums[large_index]
                    large_index -= 1
                else:
                    p += 1
            if right - large_index >= k:
                return helper(large_index + 1, right, k)
            elif right - small_index + 1 < k:
                return helper(left, small_index - 1, k - (right - small_index + 1))
            else:
                return pivot
        return helper(0, len(nums) - 1, k)

堆

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = []
        for each_num in nums:
            if len(heap) < k:
                heapq.heappush(heap, each_num)
            elif each_num > heap[0]:
                heapq.heapreplace(heap, each_num)
        return heap[0]