leetcode215. 数组中的第K个最大元素，小根堆/快排思想

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

给定整数数组 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 <= 105
-104 <= nums[i] <= 104

题目分析

给定一个整数数组 nums 和一个整数 k，返回数组中第 k 小的元素。这个问题可以通过最小堆（优先队列）/快排思想来解决。

方法一：自建小根堆

堆（优先队列）的概念介绍

堆的定义

堆是一种特殊的完全二叉树，其中每个父节点的值都必须小于或等于其所有子节点的值（小根堆）或每个父节点的值都必须大于或等于其所有子节点的值（大根堆）。

大根堆和小根堆

• 大根堆：每个父节点的值都必须大于或等于其所有子节点的值。
• 小根堆：每个父节点的值都必须小于或等于其所有子节点的值。

数组表示堆

在计算机中，堆通常使用数组来表示。堆的根节点位于数组的第一个位置，数组的索引从0开始。

建堆的过程

堆的构建过程就是从最后一个非叶子结点开始从下往上调整。 最后一个非叶子节点怎么找？这里我们用数组表示待排序序列，则最后一个非叶子结点的位置是：数组长度/2-1。

找前K个最大的数应该用小根堆

• 原因：使用小根堆可以在O(n)时间内找到第K大的数，这是因为每次从堆顶取出最小的数，直到剩下K个数，堆顶即为第K大的数。
• 实现：初始化一个大小为K的小根堆，遍历数组，将每个元素与堆顶元素比较，如果小于堆顶元素，则将该元素与堆顶元素交换，并调整堆。重复此过程，直到数组遍历结束。

算法步骤

1. 初始化一个最小堆 a，大小为 k。
2. 调用 buildMinHeap(a, k) 函数构建最小堆。
3. 从后往前遍历数组 nums，每次将堆顶元素与当前遍历到的元素比较。
4. 如果当前元素小于堆顶元素，则将当前元素与堆顶元素交换，然后对堆进行 adjMinHeap 操作。
5. 重复步骤3和4，直到数组遍历结束。
6. 堆顶元素即为第 k 小的元素。

算法流程

具体代码

class Solution {
public:
    void adjMinHeap(vector<int>& nums, int root, int heapsize) {
        int left = root * 2 + 1, right = root * 2 + 2, minimum = root;
        if (left < heapsize && nums[left] < nums[minimum])
            minimum = left;
        if (right < heapsize && nums[right] < nums[minimum])
            minimum = right;
        if (minimum != root) {
            swap(nums[minimum], nums[root]);
            adjMinHeap(nums, minimum, heapsize);
        }
    }

    void buildMinHeap(vector<int>& nums, int k) {
        for (int i = k / 2 - 1; i >= 0; i--)
            adjMinHeap(nums, i, k);
    }
    int findKthLargest(vector<int>& nums, int k) {
        buildMinHeap(nums, k);
        for (int i = k; i < nums.size(); i++) {
            if (nums[i] < nums[0])
                continue;
            swap(nums[0], nums[i]);
            adjMinHeap(nums, 0, k);
        }
        return nums[0];
    }
};

堆（优先队列）的构建与调整复杂度分析

建堆的复杂度

• 时间复杂度：O(k)，其中 k 是堆的大小。
• 原因：建堆的过程是从数组的中间位置开始，然后逐步向上调整，直到根节点。由于堆的大小是 k，因此建堆的时间复杂度是 O(k)。

调整堆的复杂度

• 时间复杂度：O(n log k)，其中 n 是数组的大小，k 是堆的大小。
• 原因：每次调整堆的时间复杂度是 O(log k)，因为堆的调整是通过递归进行的，每次递归都会将问题规模减少一半。由于需要对数组中的每个元素进行一次调整，因此总的时间复杂度是 O(n log k)。

总的时间复杂度

• 总时间复杂度：O(k + n log k)
• 为什么是 O(n)：尽管建堆和调整堆的时间复杂度分别是 O(k) 和 O(n log k)，但由于 k 的数量级通常远小于 n，因此总的时间复杂度近似为 O(n log k)。在实际应用中，通常可以忽略 k 相对于 n 的影响，从而认为总的时间复杂度是 O(n)。

空间复杂度

• O(k)，因为只需要存储 k 个元素的最大堆。

方法二：快排思想

题目分析

快速选择算法是一种用于在未排序的数组中找到第 k 大的元素的算法。它基于快速排序的思想，但与快速排序不同，快速选择算法只对部分数组进行排序，从而找到第 k 大的元素。

算法步骤

1. 选择数组的第一个元素作为划分元素。
2. 将数组分为两部分：小于划分元素的和大于划分元素的。
3. 如果 k 小于等于划分元素左侧的元素数量，则在左侧递归查找第 k 大的元素。
4. 如果 k 大于划分元素左侧的元素数量，则在右侧递归查找第 k 大的元素。
5. 重复步骤2-4，直到找到第 k 大的元素。

算法流程

具体代码

class Solution {
public:
    // 快速选择函数，用于找到数组中第k大的元素
    int quickselect(vector<int> &nums, int l, int r, int k) {
        // 如果左指针等于右指针，说明数组只有一个元素，直接返回该元素
        if (l == r)
            return nums[k];
        // 选择左边界作为划分元素
        int partition = nums[l], i = l - 1, j = r + 1;
        // 双指针遍历，将小于划分元素的放在左边，大于的放在右边
        while (i < j) {
            // 找到第一个大于等于划分元素的索引
            do i++; while (nums[i] < partition);
            // 找到第一个小于等于划分元素的索引
            do j--; while (nums[j] > partition);
            // 如果两个指针还没有相遇，交换它们指向的元素
            if (i < j)
                swap(nums[i], nums[j]);
        }
        // 如果k小于等于右指针，说明第k大的元素在左边，递归左半部分
        if (k <= j)return quickselect(nums, l, j, k);
        // 否则在右边，递归右半部分
        else return quickselect(nums, j + 1, r, k);
    }

    // 主函数，用于找到数组中第k大的元素
    int findKthLargest(vector<int> &nums, int k) {
        // 获取数组长度
        int n = nums.size();
        // 调用快速选择函数，注意第k大的元素实际上是第n-k小的元素
        return quickselect(nums, 0, n - 1, n - k);
    }
};

算法分析

• 时间复杂度: O(n)，最坏情况下时间复杂度为 O(n^2)，但平均情况下接近 O(n)。
• 空间复杂度: O(log n)，由于使用了递归，递归深度为 log n。
• 应用场景: 适用于查找第 k 大的元素，尤其当 k 远小于 n 时，效率较高。

相似题目

题目	链接
215. 数组中的第K个最大元素	https://leetcode.cn/problems/kth-largest-element-in-an-array/
218. 天际线问题	https://leetcode.cn/problems/the-skyline-problem/
234. 回文链表	https://leetcode.cn/problems/palindrome-linked-list/