「力扣」第 215 题：数组第 k 大的元素

专题 3：数组

发布日期: 2017-06-13

中文网址：215. 数组中的第K个最大元素；
英文网址：215. Kth Largest Element in an Array 。

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

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

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

思路1：排序，然后返回倒数第 $k$ 个元素，索引是 $n - k$；

思路2：partition ，逐渐减少搜索的范围，partition 的核心是大于等于的放过，小于的才做操作，因为要让小于的挪到前面去，还能保证元素的相对位置不变；注意一些边边角角的细节，+1 和 -1 要特别小心。

Python 代码：partition 的过程一定要在理解的基础上，熟记

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

    # 数组中的第 K 个最大元素
    # 数组中第 k 大的元素，它的索引是 len(nums) - k
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """

        size = len(nums)

        if size < k:
            raise Exception('程序出错')
            # [0,1,2,3,4,5]

        # 第 k 大元素的索引是 len(nums) - k
        left = 0
        right = len(nums) - 1

        while True:
            index = self.__partition(nums, left, right)
            if index == len(nums) - k:
                return nums[index]
            if index > len(nums) - k:
                right = index - 1
            else:
                left = index + 1

    def __partition(self, nums, left, right):
        """
        partition 是必须要会的子步骤，一定要非常熟练
        在 [left, right] 这个区间执行 partition
        遇到比第一个元素大的或等于的，就放过，遇到小的，就交换
        :param nums:
        :param left:
        :param right:
        :return:
        """
        pivot = nums[left]
        k = left
        for index in range(left + 1, right + 1):
            if nums[index] < pivot:
                k += 1
                nums[k], nums[index] = nums[index], nums[k]
        nums[left], nums[k] = nums[k], nums[left]
        return k


if __name__ == '__main__':
    nums = [3, 7, 8, 1, 2, 4]
    solution = Solution()
    result = solution.findKthLargest(nums, 2)
    print(result)

思路3：使用堆。

Python 代码1：使用容量为 k 的小顶堆，元素个数小于 k 的时候，放进去就是了；元素个数大于 k 的时候，小于堆顶元素，就扔掉，大于堆顶元素，就替换。

import heapq


class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """

        size = len(nums)
        if k > size:
            raise Exception('程序出错')

        # 堆有序数组
        h = []

        for num in nums:
            if len(h) < k:
                heapq.heappush(h, num)
            else:
                if num < h[0]:
                    pass
                else:
                    heapq.heappushpop(h, num)
        return h[0]

Python 代码2：与 Python 代码1 等价的写法

import heapq


# 还可以参考：https://leetcode.com/problems/kth-largest-element-in-an-array/discuss/167837/Python-or-tm

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        L = []
        for index in range(k):
            # 默认是最小堆
            heapq.heappush(L, nums[index])
        for index in range(k, len(nums)):
            top = L[0]
            if nums[index] > top:
                # 看一看堆顶的元素，只要比堆顶元素大，就替换堆顶元素
                heapq.heapreplace(L, nums[index])
        # 最后堆顶中的元素就是堆中最小的，整个数组中的第 k 大元素
        return L[0]

Python 代码3：使用大顶堆，全部放进去以后，再往外 pop

import heapq


class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        l = [(-num, num) for num in nums]
        heapq.heapify(l)
        for _ in range(k - 1):
            heapq.heappop(l)
        return l[0][1]

说明：Python 中的 heapq 可以传入 tuple，heapq 会根据 tuple 的 0 号索引元素进行堆的操作。

Python 代码4：与 Python 代码 3 等价的写法

import heapq


class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        l = [(-num, num) for num in nums]
        heapq.heapify(l)
        for _ in range(k):
            _, res = heapq.heappop(l)
        return res

（本节完）

题解：通过 partition 减治 + 优先队列（Java、C++、Python）

传送门：215. 数组中的第K个最大元素。

传送门：英文网址：215. Kth Largest Element in an Array ，中文网址：215. 数组中的第K个最大元素。

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

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

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

这道题应该说是无比重要的高频考题，是一定要掌握的。

两种思路分别使用了很基础的数据结构（优先队列）和算法（partition）。

方法一：使用快速排序 partition 的思路

Python 代码：

class Solution:

    # 数组中的第 K 个最大元素
    # 数组中第 k 大的元素，它的索引是 len(nums) - k
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """

        left = 0
        right = len(nums) - 1

        while True:
            index = self.__partition(nums, left, right)
            if index == len(nums) - k:
                return nums[index]
            if index > len(nums) - k:
                right = index - 1
            else:
                left = index + 1

    def __partition(self, nums, left, right):
        """
        partition 是必须要会的子步骤，一定要非常熟练
        典型的例子就是：[3,7,8,1,2,4]
        遇到比第一个元素大的或等于的，就放过，遇到小的，就交换
        在 [left,right] 这个区间执行 partition
        :param nums:
        :param left:
        :param right:
        :return:
        """
        pivot = nums[left]
        k = left
        for index in range(left + 1, right + 1):
            if nums[index] < pivot:
                k += 1
                nums[k], nums[index] = nums[index], nums[k]
        nums[left], nums[k] = nums[k], nums[left]
        return k


if __name__ == '__main__':
    nums = [3, 7, 8, 1, 2, 4]
    solution = Solution()
    result = solution.findKthLargest(nums, 2)
    print(result)

解法2：使用优先队列