数组中的第K个最大元素 -- 堆&快排

题目描述

LeetCode链接：数组中的第K个最大元素

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

class FindKthLargest:
    """
    https://leetcode.cn/problems/kth-largest-element-in-an-array/
    215. 数组中的第K个最大元素
    """

    def solution1(self, k: int, nums: List[int]):
        """
        小根堆
        构建一个大小为K的小根堆，放入K个元素。遍历元素，大于堆顶加入堆，维护堆的大小为k。
        最后堆顶即可第k大元素。
        时间复杂度：O(N*logK)
        空间复杂度：O(K)
        :param k:
        :param nums:
        :return:
        """
        min_heap = []
        for num in nums:
            heapq.heappush(min_heap, num)
            if len(min_heap) > k:
                heapq.heappop(min_heap)
        return min_heap[0]

    def solution2(self, k: int, nums: List[int]):
        """
        大根堆
        所有元素都放入堆，最后pop()k次即可。
        时间复杂度：O(N*logN)
        空间复杂度：O(N)
        :param k:
        :param nums:
        :return:
        """
        max_heap = []
        # 所有元素都放入堆
        for num in nums:
            heapq.heappush(max_heap, -num)

        res = None
        for _ in range(k):
            res = -heapq.heappop(max_heap)
        return res

    def solution3(self, k: int, nums: List[int]):
        """
        快排
        时间复杂度：O(N)
        空间复杂度：O(N)
        :return:
        """
        return self.findK(nums, k, 0, len(nums)-1)

    def findK(self, nums, k, low, high):
        if low == high:
            return nums[low]

        partition = self.quickSortOneTime(nums, low, high)
        # 找到的是第几个大值
        idx = high - partition + 1
        if idx == k:
            return nums[partition]
        elif idx < k:
            return self.findK(nums, k-idx, low, partition-1)
        else:
            return self.findK(nums, k, partition+1, high)

    def quickSortOneTime(self, nums, low, high):
        """
        一趟快速排序，即【快速选择】
        :param nums:
        :param low:
        :param high:
        :return:
        """
        k_value = nums[low]
        while low < high:
            while k_value < nums[high] and low < high:
                high -= 1
            nums[low] = nums[high]

            while k_value >= nums[low] and low < high:
                low += 1
            nums[high] = nums[low]

        nums[high] = k_value
        return high