侯瑞哲的博客

Ruizhe's Blog

0%

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

发表于 Valine：

借助本题深入了解快速排序中的随机选择，和堆的堆化、插入、删除操作

题目

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

在未排序的数组中找到第 k 个最大的元素。请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

示例 1:

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

示例 2:

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

说明:

你可以假设 k 总是有效的，且 1 ≤ k ≤ 数组的长度。

思路

  1. 暴力解法

题目要求我们找到”数组排序后的第k个最大的元素，而不是第k个不同的元素“。

因此，升序排序以后，目标元素的索引是 len - k ，这是最简单的思路。

  • 最简单同时也一定是最容易编码的，编码成功的几率最高，可以用这个简单思路编码的结果和其他思路编码的结果进行比对，验证高级算法的正确性；
  • 在数据规模小、对时间复杂度、空间复杂度要求不高的时候，简单问题简单做；
  • 思路简单的算法考虑清楚了，有些时候能为实现高级算法铺路，这道题也是如此；
  • 低级算法往往容错性最好，即在输入不满足题目条件的时候，往往还能得到正确的答案，而高级算法对输入数据的要求就非常苛刻。

复杂度分析：

  • 时间复杂度：O(nlogn)，这里n是数组的长度，主要是排序，语言内置的排序一般是优化的快排或者归并，O(nlogn)
  • 空间复杂度：O(1)

  1. 借助 partition 操作定位到最终排定以后索引为 len - k 的那个元素（特别注意：随机化切分元素）

以下是注意事项，因为很重要，所以放在前面说：

快速排序虽然快，但是如果实现得不好，在遇到特殊测试用例的时候，时间复杂度会变得很高。

分析：我们在学习"快速排序"的时候，接触的第1个操作就是partition（切分），简单介绍如下：

partition（切分）操作，使得：

  • 对于某个索引j，nums[j]已经排定，即nums[j]经过partition（切分）操作以后会放置在它“最终应该放置的地方”；
  • nums[left]到nums[j - 1]中的所有元素都不大于nums[j]；
  • nums[j + 1]到nums[right]中的所有元素都不小于nums[j]；

partition（切分）操作总能排定一个元素，还能够知道这个元素它最终所在的位置，这样每经过一次partition（切分）操作就能缩小搜索的范围，这样的思想叫做“减而治之”（是分而治之的思想的特例）。

切分过程可以不借助额外的数组空间，仅通过交换数组元素来实现。

1
2
3
4
5
6
7
8
9
def partition(nums, l, r):
    pivot = r
    right = l
    for i in range(l, r):
        if nums[i] <= nums[pivot]:
            nums[i], nums[right] = nums[right], nums[i]
            right += 1
    nums[right], nums[pivot] = nums[pivot], nums[right]
    return right

复杂度分析：

  • 时间复杂度：O(N)，这里N是数组的长度，可以通过主定理证明。
  • 空间复杂度：O(1)

注意：本地必须随机初始化pivot元素，否则通过时间会很慢，因为测试中有极端用例。

  1. 为了应对极端测试用例，使得递归树加深，可以在循环一开始的时候，随机交换第1个元素与它后面的任意1个元素的位置；

说明：最极端的是顺序数组与倒序数组，此时递归树画出来是链表，时间复杂度退化为O(N^2)，根本达不到减治的效果。

1
2
3
4
5
6
7
8
9
10
11
def partition(nums, l, r):
    rand = random.randint(l, r)
    nums[r], nums[rand] = nums[rand], nums[r]
    pivot = r
    right = l
    for i in range(l, r):
        if nums[i] <= nums[pivot]:
            nums[i], nums[right] = nums[right], nums[i]
            right += 1
    nums[right], nums[pivot] = nums[pivot], nums[right]
    return right
  1. 使用双指针，将与pivot相等的元素等概率地分到pivot最终排定位置的两边。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def partition(nums, l, r):
    # randint 是包括左右区间的
    random_index = random.randint(l, r)
    nums[random_index], nums[l] = nums[l], nums[random_index]

    pivot = nums[l]

    lt = l + 1
    rt = r

    while True:
        while lt <= rt and nums[lt] < pivot:
            lt += 1
        while lt <= rt and nums[rt] > pivot:
            rt -= 1

        if lt > rt:
            break
        nums[lt], nums[rt] = nums[rt], nums[lt]
        lt += 1
        rt -= 1

    nums[l], nums[rt] = nums[rt], nums[l]
    return rt

完整快速选择实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        def partition(nums, l, r):
            pivot = r
            rand = random.randint(l, r)
            nums[pivot], nums[rand] = nums[rand], nums[pivot]
            right = l
            for i in range(l, r):
                if nums[i] <= nums[pivot]:
                    nums[i], nums[right] = nums[right], nums[i]
                    right += 1
            nums[right], nums[pivot] = nums[pivot], nums[right]
            return right

        l, r = 0, len(nums) - 1
        while True:
            pivot = partition(nums, l, r)
            if pivot == len(nums) - k:
                return nums[pivot]
            if pivot < len(nums) - k:
                l = pivot + 1
            else:
                r = pivot - 1
  1. 优先队列、堆

优先队列的想法是很朴素的。

我们可以维护大小为K的一个小顶堆：

  1. 如果当前堆不满，直接添加；
  2. 堆满的时候，如果新读到的数小于等于堆顶，肯定不是我们要找的元素，大于堆顶，我们pop堆顶元素，插入新到的元素，让堆自己去调整内部结构。

说明：这里最合适的操作其实是replace，即直接把新读进来的数放在堆顶，然后执行下沉操作（siftDown）操作。python中有heapq.heappushpop()操作。Java中的PriorityQueue没有提供这个操作，只好先poll()再offer()。

堆的写法有很多，以下的写法大同小异，没有本质差别。

假设数组有len个元素。

思路1：把len个元素都放入一个小顶堆堆中，然后在pop()出len-k个元素，此时小顶堆只剩下k个元素，堆顶元素就是数组的第k大元素。

思路2：把len个元素都放入一个大顶堆中，然后再pop出k-1个元素，因此前k-1大的元素都被弹出了，此时大顶堆的堆顶元素就是数组中的第k个最大元素。

下面给出第2中思路的实现，这里改变了原数组，如果不能改变原数组，我们要另外使用堆的空间，可以为k规模的大小，也可以为len规模的大小。

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
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        def maxheap(a, i, size):
            l = i * 2 + 1
            r = i * 2 + 2
            large = i
            if l < size and a[l] > a[large]:
                large = l
            if r < size and a[r] > a[large]:
                large = r
            if large != i:
                nums[i], nums[large] = nums[large], nums[i]
                maxheap(a, large, size)

        def buildheap(a, size):
            for i in range(size >> 1, -1, -1):
                maxheap(a, i, size)

        heapsize = len(nums)
        buildheap(nums, heapsize)
        for i in range(heapsize - 1, heapsize - k, -1):
            nums[0], nums[i] = nums[i], nums[0]
            heapsize -= 1
            maxheap(nums, 0, heapsize)
        return nums[0]

考察点

这道题目考查了快速选择（快速切分）操作的特性，时间复杂度的了解，和随机化保证时间复杂度不会退化得过于严重。

堆的解法考查了堆化、插入、和删除的代码写法，是非常好的面试题目。

参考

https://leetcode-cn.com/problems/kth-largest-element-in-an-array/solution/partitionfen-er-zhi-zhi-you-xian-dui-lie-java-dai-/