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

题目说明

题目链接：
https://leetcode.cn/problems/kth-largest-element-in-an-array/

给定整数数组 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

题目解答

题目要求再$O(N)$的时间复杂度下完成，所以不能直接对数组进行排序操作，因为排序的时间复杂度为$O(NlogN)$。对于这个问题，有两种解决方案，其中最简单的一种方法是堆（优先队列），代码如下：

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int,vector<int>,greater<int> > pq;
        for(int i=0;i<nums.size();i++){
            if(pq.size()<k)pq.push(nums[i]);
            else if(pq.top()<nums[i]){
                pq.pop();
                pq.push(nums[i]);
            }   
        }
        return pq.top();
    }
};

创建一个大小为$k$的优先队列来模拟堆，该堆是一个小顶堆，用来存放数组中最大的$k$个元素。当优先队列的元素小于$k$时，直接将元素放入优先队列之中；如果优先队列的元素已经为$k$，那么比较堆顶元素的大小，如果比堆顶大则插入，同时把原堆顶弹出。该算法时间复杂度仍为$O(NlogN)$，空间复杂度为$O(k)$。还有一种使用堆的思路是原地建立一个大小为$N$的大顶堆，做$k-1$次删除操作即可，这种情况下的时间复杂度仍为$O(NlogN)$，空间复杂度为建堆时用的栈空间$O(logN)$。两种堆的方法差距就在空间复杂度上，一般情况下第二种的表现会更好一些。

题目的第二种解决方案是类似快速排序的策略，首先回顾一下快速排序。快速排序是一种原地排序算法，主要使用的是分治策略，将一个大的排序任务拆分为若干小的排序任务，其平均时间复杂度为$O(NlogN)$。

void quick_sort(vector<int>&nums,int l,int r){
    if(l>=r)return;
    int pl=l,pr=r,root=nums[l];
    while(pl<pr){
        // 测试过无论是大于或者大于等于的效率都一样
        // 应该是不使用等于的效率会高一些，减少比较次数
        while(pl<pr&&nums[pr]>root)pr--;
        while(pl<pr&&nums[pl]<root)pl++;
        // 交换两个指针所指向的值
        swap(nums[pr],nums[pl]);
        // 另写法不需要swap：
        // while(pl<pr&&nums[pr]>root)pr--;
        // nums[pl]=nums[pr];
        // while(pl<pr&&nums[pl]<root)pl++;
        // nums[pr]=nums[pl];
    }
    // 最后的pl和pr重合，所以使用pl或pr都一样
    nums[pl]=root;
    quick_sort(nums,l,pl-1);
    quick_sort(nums,pl+1,r);
}

那么这道题目可以借鉴快速排序的思想，快速排序是将排序$(l,r)$的问题拆分成$(l,root-1)$和$(root+1,r)$两个子问题。对于本题来说，寻找$(l,r)$上的第$k$大元素，其第$k$大元素一定在$(l,root-1)$和$(root+1,r)$两者之一，这样只需要对其中的一个子问题继续研究即可，所以该方法比纯粹的快速排序更优。

int quick_sort(vector<int>& nums,int l,int r,int k){
        int root=nums[l];
        int pl=l,pr=r;
        while(pl<pr){
            while(pl<pr&&nums[pr]<root)pr--;
            nums[pl]=nums[pr];
            while(pl<pr&&nums[pl]>=root)pl++;
            nums[pr]=nums[pl];
        }
        nums[pl]=root;
        if(pl-l==k)return root;
        else if(pl-l>k)return quick_sort(nums,l,pl-1,k);
        else return quick_sort(nums, pl+1, r, k-(pl-l+1));
    }

该算法的时间复杂度为$O(N)$，是一个线性的算法，具体算法时间复杂度的推导可以参考《算法导论》。该算法只在快速排序时占用了$O(NlogN)$空间，所以空间复杂度为$O(NlogN)$。