215. 数组中的第K个最大元素（快速排序、堆排序、Python）

学习快速排序和堆排序的应用

题目描述

方法一：快速排序

参考

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        self.quickSort(nums, 0, len(nums)-1)
        return nums[-k]

    def quickSort(self, nums, left, right):
        if left < right:
            partitionIndex = self.partition(nums, left, right)
            self.quickSort(nums, left, partitionIndex-1)
            self.quickSort(nums, partitionIndex+1, right)
        
    def partition(self, nums, left, right):
        pivot = left
        index = pivot + 1
        for i in range(index, right+1):
            if nums[i] < nums[pivot]:
                nums[index], nums[i] = nums[i], nums[index]
                index += 1
        # swap index-1 and pivot
        nums[pivot], nums[index-1] = nums[index-1], nums[pivot]
        return index - 1

复杂度分析

• 时间复杂度：$O(n\log n)$
• 空间复杂度：$O(\log n)$，递归使用栈空间的空间代价的期望为 $O(\log n)$。

方法二：基于快速排序的选择方法

我们可以用快速排序来解决这个问题，先对原数组排序，再返回倒数第 $k$ 个位置，这样平均时间复杂度是 $O(n\log n)$，但其实我们可以做的更快。

首先我们来回顾一下快速排序，这是一个典型的分治算法。我们对数组 $a[l\cdots r]$ 做快速排序的过程是（参考《算法导论》）：

• 分解： 将数组 $a[l\cdots r]a[l\cdots r]$ 「划分」成两个子数组 $a[l\cdots q-1]$、$a[q+1\cdots r]$，使得 $a[l\cdots q-1]$ 中的每个元素小于等于 $a[q]$，且 $a[q]$ 小于等于 $a[q+1\cdots r]$ 中的每个元素。其中，计算下标 $q$ 也是「划分」过程的一部分。
• 解决： 通过递归调用快速排序，对子数组 $a[l\cdots q-1]$ 和 $a[q+1\cdots r]$ 进行排序。
• 合并： 因为子数组都是原址排序的，所以不需要进行合并操作，$a[l\cdots r]$ 已经有序。
• 上文中提到的 「划分」 过程是：从子数组 $a[l\cdots r]$ 中选择任意一个元素 $x$ 作为主元，调整子数组的元素使得左边的元素都小于等于它，右边的元素都大于等于它， $x$ 的最终位置就是 $q$。

由此可以发现每次经过「划分」操作后，我们一定可以确定一个元素的最终位置，即 $x$ 的最终位置为 $q$，并且保证 $a[l\cdots q-1]$ 中的每个元素小于等于 $a[q]$，且 $a[q]$ 小于等于 $a[q+1\cdots r]$ 中的每个元素。所以只要某次划分的 $q$ 为倒数第 $k$ 个下标的时候，我们就已经找到了答案。 我们只关心这一点，至于 $a[l\cdots q-1]$ 和 $a[q+1\cdots r]$ 是否是有序的，我们不关心。

因此我们可以改进快速排序算法来解决这个问题：在分解的过程当中，我们会对子数组进行划分，如果划分得到的 $q$ 正好就是我们需要的下标，就直接返回 $a[q]$；否则，如果 $q$ 比目标下标小，就递归右子区间，否则递归左子区间。 这样就可以把原来递归两个区间变成只递归一个区间，提高了时间效率。这就是 「快速选择」 算法。

我们知道快速排序的性能和「划分」出的子数组的长度密切相关。直观地理解如果每次规模为 $n$ 的问题我们都划分成 $1$ 和 $n-1$，每次递归的时候又向 $n-1$ 的集合中递归，这种情况是最坏的，时间代价是 $O(n^2)$。我们可以引入随机化来加速这个过程，它的时间代价的期望是 $O(n)$， 证明过程可以参考「《算法导论》9.2：期望为线性的选择算法」。

import random
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return self.quickSelect(nums, 0, len(nums)-1, len(nums)-k)

    def quickSelect(self, nums, left, right, index):
        partitionIndex = self.partition(nums, left, right)
        if partitionIndex == index:
            return nums[partitionIndex]
        elif partitionIndex > index:
            return self.quickSelect(nums, left, partitionIndex-1, index)
        else:
            return self.quickSelect(nums, partitionIndex+1, right, index)
        
    def partition(self, nums, left, right):
        pivot = random.randint(left, right)     # 随机选择pivot，避免最坏情况
        nums[pivot], nums[left] = nums[left], nums[pivot]     # 交换pivot到最左边，方便与数组所有元素比较
        index = left + 1
        for i in range(index, right+1):
            if nums[i] < nums[left]:
                nums[index], nums[i] = nums[i], nums[index]
                index += 1
        # swap index-1 and left
        nums[left], nums[index-1] = nums[index-1], nums[left]
        return index - 1

复杂度分析

• 时间复杂度：$O(n)$，如上文所述，证明过程可以参考「《算法导论》9.2：期望为线性的选择算法」。
• 空间复杂度：$O(\log n)$，递归使用栈空间的空间代价的期望为 $O(\log n)$。

方法三：基于堆排序的选择方法

我们也可以使用堆排序来解决这个问题——建立一个大根堆，做 $k-1$次删除操作后堆顶元素就是我们要找的答案。 在很多语言中，都有优先队列或者堆的的容器可以直接使用，但是在面试中，面试官更倾向于让更面试者自己实现一个堆。所以建议读者掌握这里大根堆的实现方法，在这道题中尤其要搞懂「建堆」、「调整」和「删除」的过程。

友情提醒：「堆排」在很多大公司的面试中都很常见，不了解的同学建议参考《算法导论》或者大家的数据结构教材，一定要学会这个知识点哦! ^_

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        global arrLen
        arrLen = len(nums)
        self.buildMaxHeap(nums)  # 构造大顶堆
        for i in range(len(nums)-1, len(nums)-k, -1):
            self.swap(nums, 0, i)
            arrLen -= 1
            self.heapify(nums, 0)
        return nums[0]

    def buildMaxHeap(self, arr):
        for i in range(arrLen // 2, -1, -1):    # i为非叶节点
            self.heapify(arr, i)

    def heapify(self, arr, i):
        left = 2 * i + 1
        right = 2 * i + 2
        largest = i
        if left < arrLen and arr[left] > arr[largest]:
            largest = left
        if right < arrLen and arr[right] > arr[largest]:
            largest = right
        if largest != i:
            self.swap(arr, largest, i)
            self.heapify(arr, largest)  # 继续向下构造大顶堆

    def swap(self, arr, i, j):
        arr[i], arr[j] = arr[j], arr[i]

复杂度分析

• 时间复杂度：$O(n\log n)$，建堆的时间代价是 $O(n)$，删除的总代价是 $O(k\log n)$，因为 $k<n$，故渐进时间复杂为 $O(n+k\log n)=O(n\log n)$。
• 空间复杂度：$O(\log n)$，即递归使用栈空间的空间代价。