0215. 数组中的第K个最大元素

0215. 数组中的第K个最大元素 #

  • 标签：数组、堆排序
  • 难度：中等

题目大意 #

描述：给定一个未排序的整数数组 nums 和一个整数 k。

要求：返回数组中第 k 个最大的元素。

说明：

  • 要求使用时间复杂度为 $O(n)$ 的算法解决此问题。
  • $1 \le k \le nums.length \le 10^5$。
  • $-10^4 \le nums[i] \le 10^4$。

示例：

1
2
3
4
5
6

输入: [3,2,1,5,6,4], k = 2
输出: 5


输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

解题思路 #

很不错的一道题，面试常考。

直接可以想到的思路是：排序后输出数组上对应第 k 位大的数。所以问题关键在于排序方法的复杂度。

冒泡排序、选择排序、插入排序时间复杂度 $O(n^2)$ 太高了，很容易超时。

可考虑堆排序、归并排序、快速排序。

这道题的要求是找到第 k 大的元素，使用归并排序只有到最后排序完毕才能返回第 k 大的数。而堆排序每次排序之后，就会确定一个元素的准确排名，同理快速排序也是如此。

思路 1：堆排序 #

升序堆排序的思路如下：

  1. 将无序序列构造成第 1 个大顶堆（初始堆），使得 n 个元素的最大值处于序列的第 1 个位置。

  2. 调整堆：交换序列的第 1 个元素（最大值元素）与第 n 个元素的位置。将序列前 n - 1 个元素组成的子序列调整成一个新的大顶堆，使得 n - 1 个元素的最大值处于序列第 1 个位置，从而得到第 2 个最大值元素。

  3. 调整堆：交换子序列的第 1 个元素（最大值元素）与第 n - 1 个元素的位置。将序列前 n - 2 个元素组成的子序列调整成一个新的大顶堆，使得 n - 2 个元素的最大值处于序列第 1 个位置，从而得到第 3 个最大值元素。

  4. 依次类推，不断交换子序列的第 1 个元素（最大值元素）与当前子序列最后一个元素位置，并将其调整成新的大顶堆。直到获取第 k 个最大值元素为止。

思路 1：代码 # 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # 调整为大顶堆
        def heapify(nums, index, end):
            left = index * 2 + 1
            right = left + 1
            while left <= end:
                # 当前节点为非叶子节点
                max_index = index
                if nums[left] > nums[max_index]:
                    max_index = left
                if right <= end and nums[right] > nums[max_index]:
                    max_index = right
                if index == max_index:
                    # 如果不用交换，则说明已经交换结束
                    break
                nums[index], nums[max_index] = nums[max_index], nums[index]
                # 继续调整子树
                index = max_index
                left = index * 2 + 1
                right = left + 1
                
        # 初始化大顶堆
        def buildMaxHeap(nums):
            size = len(nums)
            # (size-2) // 2 是最后一个非叶节点，叶节点不用调整
            for i in range((size - 2) // 2, -1, -1):
                heapify(nums, i, size - 1)
            return nums

        buildMaxHeap(nums)
        size = len(nums)
        for i in range(k-1):
            nums[0], nums[size-i-1] = nums[size-i-1], nums[0]
            heapify(nums, 0, size-i-2)
        return nums[0]

思路 1：复杂度分析 #

  • 时间复杂度：$O(n \times \log_2n)$。
  • 空间复杂度：$O(1)$。

思路 2：快速排序 #

使用快速排序在每次调整时，都会确定一个元素的最终位置，且以该元素为界限，将数组分成了左右两个子数组，左子数组中的元素都比该元素小，右子树组中的元素都比该元素大。

这样，只要某次划分的元素恰好是第 k 个下标就找到了答案。并且我们只需关注第 k 个最大元素所在区间的排序情况，与第 k 个最大元素无关的区间排序都可以忽略。这样进一步减少了执行步骤。

思路 2：代码 # 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

import random

class Solution:
    # 从 arr[low: high + 1] 中随机挑选一个基准数，并进行移动排序
    def randomPartition(self, arr: [int], low: int, high: int):
        # 随机挑选一个基准数
        i = random.randint(low, high)
        # 将基准数与最低位互换
        arr[i], arr[low] = arr[low], arr[i]
        # 以最低位为基准数，然后将序列中比基准数大的元素移动到基准数右侧，比他小的元素移动到基准数左侧。最后将基准数放到正确位置上
        return self.partition(arr, low, high)
    
    # 以最低位为基准数，然后将序列中比基准数大的元素移动到基准数右侧，比他小的元素移动到基准数左侧。最后将基准数放到正确位置上
    def partition(self, arr: [int], low: int, high: int):
        pivot = arr[low]            # 以第 1 为为基准数
        i = low + 1                 # 从基准数后 1 位开始遍历，保证位置 i 之前的元素都小于基准数
        
        for j in range(i, high + 1):
            # 发现一个小于基准数的元素
            if arr[j] < pivot:
                # 将小于基准数的元素 arr[j] 与当前 arr[i] 进行换位，保证位置 i 之前的元素都小于基准数
                arr[i], arr[j] = arr[j], arr[i]
                # i 之前的元素都小于基准数，所以 i 向右移动一位
                i += 1
        # 将基准节点放到正确位置上
        arr[i - 1], arr[low] = arr[low], arr[i - 1]
        # 返回基准数位置
        return i - 1

    def quickSort(self, arr, low, high, k):
        size = len(arr)
        if low < high:
            # 按照基准数的位置，将序列划分为左右两个子序列
            pi = self.randomPartition(arr, low, high)
            if pi == size - k:
                return arr[size - k]
            if pi > size - k:
                # 对左子序列进行递归快速排序
                self.quickSort(arr, low, pi - 1, k)
            if pi < size - k:
                # 对右子序列进行递归快速排序
                self.quickSort(arr, pi + 1, high, k)

        return arr[size - k]

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

思路 2：复杂度分析 #

  • 时间复杂度：$O(n)$。证明过程可参考「算法导论 9.2：期望为线性的选择算法」。
  • 空间复杂度：$O(\log_2 n)$。递归使用栈空间的空间代价期望为 $O(\log_2n)$。

思路 3：借用标准库（不建议） #

提交代码中的最快代码是调用了 Python 的 sort 方法。这种做法适合在打算法竞赛的时候节省时间，日常练习不建议不建议这样做。

思路 3：代码 # 

1
2
3
4

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

思路 3：复杂度分析 #

  • 时间复杂度：$O(n \times \log_2n)$。
  • 空间复杂度：$O(1)$。

思路 4：优先队列 #

  1. 遍历数组元素，对于挡圈元素 num：
    1. 如果优先队列中的元素个数小于 k 个，则将当前元素 num 放入优先队列中。
    2. 如果优先队列中的元素个数大于等于 k 个，并且当前元素 num 大于优先队列的队头元素，则弹出队头元素，并将当前元素 num 插入到优先队列中。
  2. 遍历完，此时优先队列的队头元素就是第K个最大元素，将其弹出并返回即可。

这里我们借助了 Python 中的 heapq 模块实现优先队列算法，这一步也可以通过手写堆的方式实现优先队列。

思路 4：代码 # 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

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

思路 4：复杂度分析 #

  • 时间复杂度：$O(n \times \log_2k)$。
  • 空间复杂度：$O(k)$。
本站总访问量  次  |  您是本站第  位访问者