在问 · 算法题解

215. 第 K 个最大元素

中等栈与堆大小为 k 的小根堆力扣原题 ↗

题目描述

输入整数数组 nums 和整数 k。把数组从大到小排好,返回排在第 k 位的元素。

重复的值要分别计数,所以要找的不是第 k 个互不相同的值。

你需要设计时间复杂度为 O(n) 的算法。

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

解释:从大到小为 6、5、4…,第 2 大是 5。

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

解释:从大到小为 6、5、5、4…,第 4 大是 4。

数据范围

  • 1 <= k <= nums.length <= 10^5
  • -10^4 <= nums[i] <= 10^4

解题思路:大小为 k 的小根堆

维护容量为 k 的最小堆。

为什么正确

堆中保存当前最大的 k 个元素。

复杂度分析

时间 O(n log k),额外空间 O(k)。

易错点

注意边界条件并在回溯或弹出后恢复状态。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。

import heapq
class Solution:
 def findKthLargest(self,a,k):
  h=[]
  for x in a:heapq.heappush(h,x);heapq.heappop(h) if len(h)>k else None
  return h[0]

载入编辑器,自己跑一遍 →

其他解法

堆解法稳定 O(n log k),不修改原数组,也适合数据流;快速选择每次只递归一侧,平均 O(n),但最坏会退化,并且需要一次性拿到整个数组。

解法 2:快速选择 时间 平均 O(n),最坏 O(n²),空间 O(1)

第 k 大就是升序排列后下标为 n-k 的元素。借用快速排序的划分:取区间中点的值作为基准,用 Hoare 划分把区间分成左边都 ≤ 基准、右边都 ≥ 基准的两段,分界为 j。目标下标在哪一段,就只在那一段继续划分,另一段直接丢弃,直到区间只剩一个元素。

平均线性时间、原地进行,比堆快。最坏情况退化到 O(n²);会修改原数组;划分的边界条件容易写错。Hoare 划分在大量重复元素时仍能均匀划分。

from typing import List

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        target = len(nums) - k
        l, r = 0, len(nums) - 1
        while l < r:
            pivot = nums[(l + r) // 2]
            i, j = l - 1, r + 1
            while True:
                i += 1
                while nums[i] < pivot:
                    i += 1
                j -= 1
                while nums[j] > pivot:
                    j -= 1
                if i >= j:
                    break
                nums[i], nums[j] = nums[j], nums[i]
            if target <= j:
                r = j
            else:
                l = j + 1
        return nums[l]

栈与堆 专题的其他题目

查看 栈与堆 全部题目 →