LeetCode215: 数组中的第K个最大元素 —— 从快速选择到堆排

题目描述

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

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

示例 1：

text

复制代码

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

示例 2：

text

复制代码

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

提示：

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

解法一：快速选择（Quick Select）

思想

快速选择是快速排序的变种。核心步骤：

随机选择一个基准元素 pivot。
将数组划分为三部分：大于 pivot、等于 pivot、小于 pivot。
根据各部分的大小与 k 的关系，确定第 k 大的元素在哪一部分，并递归处理该部分。

由于每次只需处理一侧，平均时间复杂度为 O(n)，且可原地分区优化空间。

代码实现（Java，非原地分区，易于理解）

java 复制代码

public class Solution {
    private int quickSelect(List<Integer> nums, int k) {
        // 随机选择基准数
        Random rand = new Random();
        int pivot = nums.get(rand.nextInt(nums.size()));
        // 将大于、小于、等于 pivot 的元素划分至 big, small, equal 中
        List<Integer> big = new ArrayList<>();
        List<Integer> equal = new ArrayList<>();
        List<Integer> small = new ArrayList<>();
        for (int num : nums) {
            if (num > pivot)
                big.add(num);
            else if (num < pivot)
                small.add(num);
            else
                equal.add(num);
        }
        // 第 k 大元素在 big 中，递归划分
        if (k <= big.size())
            return quickSelect(big, k);
        // 第 k 大元素在 small 中，递归划分
        if (big.size() + equal.size() < k)
            return quickSelect(small, k - (big.size()+equal.size()));
        // 第 k 大元素在 equal 中，直接返回 pivot
        return pivot;
    }
    
    public int findKthLargest(int[] nums, int k) {
        List<Integer> numList = new ArrayList<>();
        for (int num : nums) {
            numList.add(num);
        }
        return quickSelect(numList, k);
    }
}

原地分区优化版本（推荐）

java 复制代码

public class Solution {
    private Random rand = new Random();
    
    private int quickSelect(int[] nums, int left, int right, int k) {
        int pivotIdx = left + rand.nextInt(right - left + 1);
        int pivot = nums[pivotIdx];
        
        // 三路分区 (荷兰国旗问题)
        int lt = left, i = left, gt = right;
        while (i <= gt) {
            if (nums[i] > pivot) {
                swap(nums, lt++, i++);
            } else if (nums[i] < pivot) {
                swap(nums, i, gt--);
            } else {
                i++;
            }
        }
        // 此时: [left, lt-1] > pivot, [lt, gt] == pivot, [gt+1, right] < pivot
        
        if (k <= lt - left) {
            return quickSelect(nums, left, lt - 1, k);
        } else if (k <= gt - left + 1) {
            return pivot;
        } else {
            return quickSelect(nums, gt + 1, right, k - (gt - left + 1));
        }
    }
    
    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
    
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, 0, nums.length - 1, k);
    }
}

复杂度分析

时间复杂度：平均 O(n)，最坏 O(n²)（但随机化后概率极低）
空间复杂度：递归栈深度 O(log n)，原地分区额外 O(1)

解法二：最小堆（大小为 k）

思想

维护一个大小为 k 的最小堆，堆顶始终是当前扫描过的元素中第 k 大的候选值。

遍历数组，若堆大小 < k，直接入堆。
否则，若当前元素 > 堆顶，则弹出堆顶并压入当前元素（说明当前元素更有资格成为第 k 大）。
遍历结束后，堆顶就是整个数组的第 k 大元素。

代码实现

java 复制代码

import java.util.PriorityQueue;

public class Solution {
    public int findKthLargest(int[] nums, int k) {
        // 最小堆
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        for (int num : nums) {
            if (minHeap.size() < k) {
                minHeap.offer(num);
            } else if (num > minHeap.peek()) {
                minHeap.poll();
                minHeap.offer(num);
            }
        }
        return minHeap.peek();
    }
}

复杂度分析

时间复杂度：O(n log k)，每个元素最多一次入堆和出堆操作。
空间复杂度：O(k)，用于存储堆。

为什么不用最大堆？

最大堆需要将所有元素入堆（O(n) 建堆），然后弹出 k-1 次，总时间 O(n + k log n)，当 k 接近 n 时性能不错，但空间 O(n) 较大。最小堆更节省内存，且对 k 较小的情况更优。

总结

快速选择：分治思想，平均线性时间，需要处理分区细节，推荐掌握。
最小堆：代码简单，适合 k 较小时，空间占用低。
排序法：适合快速原型，生产环境如果 n 不大也可以用。

面试时，建议先说出快速选择原理，然后写出堆排序作为兜底实现。两者都能通过 LeetCode 测试，重点在于分析不同解法的优劣。