【OJ - 堆】找数组中的第K个最大元素（C++ STL - priority_queue）

本文已参与「新人创作礼」活动， 一起开启掘金创作之路

OJ - 找数组中的第K个最大元素

题目难度：中等

LeetCode链接：215. 数组中的第K个最大元素

题目描述：

---

解题思路

思路1：暴力求解，排序

使用头文件 <algorithm> 中的 sort 函数模板，直接排序，找到第k个最大的元素

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        // 对数组nums排序
        // 类模板参数默认是less，升序排序 <
        // 这里传类模板参数greater，降序排序 >
        sort(nums.begin(), nums.end(), greater<int>());
		// 返回第k个最大的元素
        return nums[k - 1];
    }
};

• 时间复杂度：因为 sort 底层是快速排序，所以时间复杂度为 $O(N*logN)$
• 空间复杂度：$O(1)$

---

思路2：建大堆

使用堆排序，建立 n 个元素的大堆，做 k - 1 次删除操作后，此时堆顶元素就是第k个最大的元素

因为 C++ 的优先级队列 priority_queue 本质就是堆，所以这道题我没有自己建堆

【n 个元素的堆，插入和删除元素的时间复杂度都为 O(logN)】

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        // 建立 n 个元素的大堆，用迭代器区间[first,last)初始化
        priority_queue<int> maxHeap(nums.begin(), nums.end());
        
        /* 上述入堆操作也可以这样写：
        // 将数组nums中的元素依次放入优先级队列中
        priority_queue<int> maxHeap;  // priority_queue默认是大堆
        for(size_t i = 0; i < nums.size(); i++) {
            maxHeap.push(nums[i]);
        }
        */

        // popd堆顶元素k-1次后，此时的堆顶元素就是第k个最大的元素
        while(--k) {
            maxHeap.pop();
        }
        
        return maxHeap.top();
    }
};

• 时间复杂度：向下调整建立 n 个数的大堆时间代价为 O(N)；删除堆顶元素 k - 1 次（交换后删除，并向下调整）的时间代价为 O(klogN)；相加为 $O(N + klogN)$​
• 空间复杂度：建立n个数的大堆，需要开空间存这n个数，所以为 $O(N)$​

---

思路3：建小堆，最优解

思考：如果数组 nums 中的元素个数 n 特别大，前两种方法的弊端就出来了，那么如果做到时间和空间复杂度更优呢？

因为我们要找的是第 k 个最大的元素：

• 建立 k 个元素的小堆，数组 nums 前 k 个元素依次放入小堆，此时的堆顶元素是堆中最小的元素，也是堆里面第 k 个最小的元素，

• 然后把数组 nums 剩下的元素与堆顶比较，若大于堆顶则去掉堆顶，再将其插入，

• 这样一来，堆里面存放的就是数组 nums 中的前 k 个最大元素，

• 此时小堆的堆顶元素也就是堆中的第 k 个最大的元素

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
		// 建立 k 个元素的小堆，传类模板参数 greater，将 priority_queue 定义成小堆
        priority_queue<int, vector<int>, greater<int>> minHeap;
        
        // 将数组 nums 前 k 个元素依次放入优先级队列中
        for(size_t i = 0; i < k; i++)
        {
            minHeap.push(nums[i]);
        }

        // 遍历nums剩下的元素，与堆顶比较，若大于堆顶则去掉堆顶，再将其插入
        for(size_t i = k; i < nums.size(); i++)
        {
            if(nums[i] > minHeap.top())
            {
                minHeap.pop();          // pop掉当前堆顶元素
                minHeap.push(nums[i]);  // 入堆数组元素
            }
        }
        
        // 结束后小堆的堆顶元素就是第k个最大的元素
        return minHeap.top();
    }
};

• 时间复杂度：建立 k 个数的小堆时间代价为 O(k)；剩余 n - k 个数与堆顶元素比较一下，pop一下，push一下，时间代价为 O((N-k)*logk)；相加为 $O(k + (N-k)logk)$​

  当 n 远大于 k 时，这种方法更优，尤其是 n 大到内存中存不下。

• 空间复杂度：建立 k 个数的大堆，需要开空间存这 k 个数，所以为 $O(k)$

---