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

力扣题目链接 (opens new window)

题目描述

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

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

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:

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

示例 2:

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

提示：

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

思路

这道题有两种经典做法：小顶堆和快速选择算法。

方法一：小顶堆（优先队列）

什么是小顶堆？就是堆顶元素是最小的。

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

• 遍历数组，把元素加入堆
• 如果堆的大小超过K，就弹出堆顶（最小元素）
• 最后堆顶就是第K个最大元素

为什么？

因为我们保留了K个最大的元素，堆顶是这K个元素中最小的，也就是第K大的。

举个例子，nums = [3,2,1,5,6,4], k = 2：

遍历到3: 堆 = [3]
遍历到2: 堆 = [2, 3]
遍历到1: 堆 = [2, 3]（1比堆顶2小，不加入）
遍历到5: 堆 = [3, 5]（加入5，弹出最小的2）
遍历到6: 堆 = [5, 6]（加入6，弹出最小的3）
遍历到4: 堆 = [5, 6]（4比堆顶5小，不加入）
结果：堆顶是5

时间复杂度：O(n log k)，每次堆操作是log k，总共n个元素。

方法二：快速选择算法（Quick Select）

快速选择算法是快速排序的变体，平均时间复杂度O(n)。

快速排序每次partition后，pivot左边的元素都比它小，右边的都比它大。

如果pivot的位置正好是第K大的位置，就找到了答案。

如果pivot位置太小，说明第K大在左边，继续在左边找。 如果pivot位置太大，说明第K大在右边，继续在右边找。

这样平均只需要遍历n + n/2 + n/4 + ... ≈ 2n = O(n)次。

注意：第K个最大元素，在排序后数组中的索引是n - k。

模拟过程

以 nums = [3,2,1,5,6,4], k = 2 为例：

小顶堆维护过程：

1. 遍历到 3：堆 = [3]
2. 遍历到 2：堆 = [2, 3]（堆大小 = K = 2）
3. 遍历到 1：1 < 2（堆顶），不加入
4. 遍历到 5：加入后堆 = [2,3,5]，超过K，弹出堆顶2，堆 = [3, 5]
5. 遍历到 6：加入后弹出堆顶3，堆 = [5, 6]
6. 遍历到 4：4 < 5（堆顶），不加入

最终堆顶 = 5，即第2大元素。

原理： 维护大小为K的小顶堆，堆中保留K个最大元素，堆顶就是第K大。

解题代码

方法一：小顶堆

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        // 小顶堆，C++中用priority_queue的greater实现
        priority_queue<int, vector<int>, greater<int>> minHeap;
        
        for (int num : nums) {
            minHeap.push(num);
            if (minHeap.size() > k) {
                minHeap.pop();  // 弹出最小元素
            }
        }
        
        return minHeap.top();  // 堆顶就是第K大
    }
};

方法二：快速选择

class Solution {
private:
    int quickSelect(vector<int>& nums, int left, int right, int k) {
        if (left == right) return nums[left];
        
        // partition：选择最右边为pivot
        int pivot = nums[right];
        int i = left;
        
        // 把小于pivot的放左边，大于pivot的放右边
        for (int j = left; j < right; j++) {
            if (nums[j] < pivot) {
                swap(nums[i], nums[j]);
                i++;
            }
        }
        swap(nums[i], nums[right]);  // pivot放到中间
        
        // i就是pivot的位置
        if (i == k) {
            return nums[i];  // 找到了
        } else if (i < k) {
            return quickSelect(nums, i + 1, right, k);  // 在右边找
        } else {
            return quickSelect(nums, left, i - 1, k);   // 在左边找
        }
    }
    
public:
    int findKthLargest(vector<int>& nums, int k) {
        int n = nums.size();
        return quickSelect(nums, 0, n - 1, n - k);  // 第K大 = 第(n-k)小
    }
};

复杂度分析

方法一：小顶堆

• 时间复杂度：O(n log k)
• 空间复杂度：O(k)

方法二：快速选择

• 时间复杂度：平均O(n)，最坏O(n²)
• 空间复杂度：O(log n)，递归栈空间

其他语言

Python3

小顶堆

import heapq

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # Python的heapq默认是小顶堆
        min_heap = []
        
        for num in nums:
            heapq.heappush(min_heap, num)
            if len(min_heap) > k:
                heapq.heappop(min_heap)  # 弹出最小元素
        
        return min_heap[0]  # 堆顶就是第K大

快速选择

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        def quick_select(left, right, k):
            if left == right:
                return nums[left]
            
            # partition：选择最右边为pivot
            pivot = nums[right]
            i = left
            
            # 把小于pivot的放左边，大于pivot的放右边
            for j in range(left, right):
                if nums[j] < pivot:
                    nums[i], nums[j] = nums[j], nums[i]
                    i += 1
            nums[i], nums[right] = nums[right], nums[i]  # pivot放到中间
            
            # i就是pivot的位置
            if i == k:
                return nums[i]  # 找到了
            elif i < k:
                return quick_select(i + 1, right, k)  # 在右边找
            else:
                return quick_select(left, i - 1, k)   # 在左边找
        
        n = len(nums)
        return quick_select(0, n - 1, n - k)  # 第K大 = 第(n-k)小

Java

小顶堆

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // 小顶堆
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        
        for (int num : nums) {
            minHeap.offer(num);
            if (minHeap.size() > k) {
                minHeap.poll();  // 弹出最小元素
            }
        }
        
        return minHeap.peek();  // 堆顶就是第K大
    }
}

Go

import "container/heap"

type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }  // 小顶堆
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func findKthLargest(nums []int, k int) int {
    h := &IntHeap{}
    heap.Init(h)
    
    for _, num := range nums {
        heap.Push(h, num)
        if h.Len() > k {
            heap.Pop(h)  // 弹出最小元素
        }
    }
    
    return (*h)[0]  // 堆顶就是第K大
}

JavaScript

var findKthLargest = function(nums, k) {
    // JavaScript没有内置堆，用数组模拟小顶堆
    const minHeap = [];
    
    for (const num of nums) {
        minHeap.push(num);
        minHeap.sort((a, b) => a - b);  // 简化版，实际应该用堆操作
        
        if (minHeap.length > k) {
            minHeap.shift();  // 弹出最小元素
        }
    }
    
    return minHeap[0];  // 堆顶就是第K大
};

与代码随想录联系

这道题是堆的经典应用。

我在代码随想录堆章节里，详细讲解了堆的原理和应用，建议录友们先看代码随想录堆章节，特别是347.前K个高频元素，那道题也是用堆来找TopK。

本题的关键点：

1. 小顶堆：维护K个最大元素，堆顶是第K大
2. 快速选择：类似快速排序，但只需要处理一边，平均O(n)

另外，很多录友会问：为什么找第K大用小顶堆，而不是大顶堆？

因为我们要保留K个最大的元素，每次淘汰的是最小的那个，所以用小顶堆。

如果用大顶堆，堆顶是最大的，但我们需要的是第K大（第K个最大的），不是最大的。