【每天一道算法题6】【Python】求数组第k个最大元素

题目描述

在未排序的数组中找到第 k 个最大的元素，

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

解法一：先排序后取值

常规做法是先对数组nums从大到小排序，然后取nums[k-1]就行。
排序算法可以选择快速排序，或者使用Python内置函数sorted()，此解法较简单，就不赘述了。时间复杂度是O(NlogN)，N是nums长度。

解法二：堆

使用Python的heapq小顶堆求解。
时间复杂度O(Nlogk) N是nums长度。

def solution(nums: list) -> int:
    # 初始化一个空列表，用于创建堆
    heap = []
    for i, n in enumerate(nums):
        if i < k:   # 先往堆中push k个元素
            heapq.heappush(heap, n)
        else:
            # 从第 k + 1个元素开始，每次和堆顶（最小值）比较，如果比堆顶大，就移除堆顶，同时把n推入堆中
            # heapreplace作用其实是先pop堆顶 然后再push
            if heap[0] < n:
                heapq.heapreplace(heap, n)
    return heap[0]

解法三：快速选择法

此方法和快速排序类似。

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # 选择枢轴
        pivot = len(nums) // 2
        tmp = nums[pivot]
        nums.remove(tmp)
        # 定义两个数组，left存放大于枢轴的数字，right存放小于枢轴的数字
        left = []
        right = []
        for n in nums:
            # 大于枢轴的放左边
            if n > tmp:
                left.append(n)
            # 小于等于枢轴的放右边
            else:
                right.append(n)
        # 如果左边的列表长度刚好等于 k - 1，那么第k大元素就是枢轴
        if len(left) == k - 1:
            return tmp
        # 如果左边的列表长度大于k-1，说明第k大的元素在左边的列表，继续对left递归求解
        elif len(left) > k - 1:
            return self.findKthLargest(left, k)
        # 如果左边的列表长度小于k-1，说明第k大的元素在右边的列表，继续对left递归求解，注意此时对右边求的是第（k - len(left) - 1）大的元素
        else:
            return self.findKthLargest(right, k - len(left) - 1)