【小白用python刷Leetcode】215. 数组中的第K个最大元素

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

题目描述

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

示例：

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

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

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

解题思路

今天每日一题刷到这道题，不禁笑出声来，终于看到了熟悉的面庞。这道应该是面试必会的一道了吧，自己秋招时候被问到了大概有三四次的样子，私以为面试时候这道题地出现频率还是相当高的（注意是面试，觉得笔试可比这个难多了，唉）。也终于不是小白思路了，应该就是通常意义上的优解了，因为知名度很高，就大概说一下，反正也没人看，权当自己做笔记。

具体来说比较好的方法就是利用堆排思想，构造一个size为K的小顶堆，然后拿剩余元素去和堆顶元素比较。因为是小顶堆，堆顶一定是堆内最小的元素，所以如果发现剩余元素中存在比当前堆顶元素还大的元素，就替换掉堆顶元素，并维护这个小顶堆。当没有剩余元素时，堆顶的元素就是第K个最大元素。由于是发现了堆外元素比堆顶元素大才会替换，所以每次维护之后，堆内元素一定是所有已经过了堆的元素中最大的K个。而堆顶又是这K个中最小的元素，所以当全部元素都过堆之后，堆顶元素就是我们要找的那个。如果要输出全部最大的K个元素，则直接输出堆即可，需要注意的是这个堆并没有排序，所以如果要输出有序的话，还需要继续排序才行。

由此延伸，其实第K个最小元素也是这么解，只不过构建的是大顶堆。简单来说就是：找最大的第K个元素，构建小顶堆；找最小的第K个元素，构建大顶堆，正好是反着的。

题解代码：

def findKthLargest(self, nums: List[int], k: int) -> int:
        
        def heap_build(parent, heap):
            child = 2 * parent + 1
            while child < len(heap):
                if child + 1 < len(heap) and heap[child + 1] < heap[child]:
                    child = child + 1
                if heap[parent] <= heap[child]:
                    break
                heap[parent], heap[child] = heap[child], heap[parent]
                parent, child = child, 2 * child + 1
            return heap

        if k > len(nums):
            return None
        heap = nums[:k]
        for i in range(k, -1, -1):
            heap_build(i, heap)
        for j in range(k, len(nums)):
            if nums[j] > heap[0]:
                heap[0] = nums[j]
                heap_build(0, heap)
        return heap[0]

既然今天的题目这么简单，不如就把最大最小K个数的代码都贴上来吧。内容和上面的代码大同小异，就不细说了。细节都在代码的注释里了，应该还挺详细的。

nums = [3,2,3,1,2,4,5,5,6]
k = 3
print('原数组：', nums)
print('排序后的数组：', sorted(nums))

#第K大的数
def heap_build(parent, heap):  # 单个枝杈下沉过程
    child = 2 * parent + 1
    while child < len(heap):
        # 先在一个小枝杈上比较左右叶子节点的大小，选小的那一个
        if child + 1 < len(heap) and heap[child + 1] < heap[child]:
            child = child + 1
        if heap[parent] <= heap[child]:  # 叶子节点再和父节点比较
            break
        # 如果父节点小，直接跳出循环，如果叶子节点小，则两个节点交换
        heap[parent], heap[child] = heap[child], heap[parent]
        parent, child = child, 2 * child + 1   # 重新定义父节点和子节点
    return heap


def Find_heap_kth(array, k):
    if k > len(array):
        return None
    heap = array[:k]
    for i in range(k, -1, -1):
        # 构建大小为k的小顶堆，循环是因为每个点都要下沉
        heap_build(i, heap)
    for j in range(k, len(array)):
        if array[j] > heap[0]:  # 剩余元素和小顶堆堆顶进行比较
            heap[0] = array[j]
            heap_build(0, heap)  # 继续维护这个堆
    return heap[0], heap
    # 第K大的数和K个最大的数（但是未排序，如果要输出有序，则需要继续排序）
print('输出第{}大的数：'.format(k))
num, array = Find_heap_kth(nums, k)
print(num)
print('输出全部前{}大的数：'.format(k))
print(array)


#第K小的数
def min_heap_build(parent, heap):
    child = 2 * parent + 1
    while child < len(heap):
        if child + 1 < len(heap) and heap[child + 1] > heap[child]:
            child = child + 1
        if heap[parent] >= heap[child]:
            break
        heap[parent], heap[child] = heap[child], heap[parent]
        parent, child = child, 2 * child + 1
    return heap


def Find_min_heap_kth(array, k):
    if k > len(array):
        return None
    heap = array[:k]
    for i in range(k, -1, -1):
        min_heap_build(i, heap)
    for j in range(k, len(array)):
        if array[j] < heap[0]:
            heap[0] = array[j]
            min_heap_build(0, heap)
    return heap[0], heap

print('输出第{}小的数：'.format(k))
n, a = Find_min_heap_kth(nums, k)
print(n)
print('输出全部前{}小的数：'.format(k))
print(a)

以上就是我，一个正儿八经的小白（大神们通过看代码应该也感觉出来了），对这道题的理解，欢迎诸位指正讨论，感谢阅读。

原题链接：

https://leetcode-cn.com/problems/kth-largest-element-in-an-array/