LeetCode 347. Top K Frequent Elements

本题是 找最大top k 的类似问题，而求解最大top k的解法很经典。

将前k个元素建立一个大小为k的小顶部堆，直接建堆时间 O(k)。如果顺序加入到堆中则需要 k*O(logk)。

对于剩下的 n-k 个元素，如果元素小于堆顶元素，不用管；如果大于，替换掉堆顶元素，并重新维护最小堆的性质。代码实现起来只要push进去，如果size超过k，pop即可。

这里需要 (n-k)*O(logk)。最后堆中的元素就是最大的k个元素。

 

对于本题来说，只能顺序建堆，因此采用小顶堆时间复杂度为 O(nlogk)。

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int,int> count;
        for (int num:nums) ++count[num];
        
        auto cmp=[&](const int &a, const int &b){return count[a]>count[b];};
        priority_queue<int,vector<int>,decltype(cmp)> q(cmp);
        for (auto x:count){
            int cur=x.first;
            q.push(cur);
            if (q.size()>k) q.pop();
        }
        vector<int> res;
        while (!q.empty()){
            res.push_back(q.top()); q.pop();
        }
        reverse(res.begin(),res.end());
        return res;
    }
};

 

类似考虑，如果建立大顶堆，时间复杂度为 O(nlog(n-k))，代码如下所示。

class Solution {
public:
    struct node{
        int num;
        int cnt;
        node(int n, int c):num(n),cnt(c){};
        bool operator<(const node &x) const {return cnt<x.cnt;}
    };
    
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int,int> count;
        for (int num:nums) ++count[num];
        
        vector<int> res;
        priority_queue<node> q;
        for (auto x:count){
            q.push(node(x.first,x.second));
            if (q.size()>count.size()-k){
                res.push_back(q.top().num);
                q.pop();
            }   
        }

        return res;
    }
};