Tag: Top-K

好久不刷题了，今晚来一道。

题目

Given a non-empty array of integers, return the k most frequent elements.
For example,
Given [1,1,1,2,2,3] and k = 2, return [1,2].
Note:

• You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
• Your algorithm’s time complexity must be better than O(n log n), where n is the array’s size.

思路

题目中明确提示要O(nlogn)复杂度了…
我的做法是先统计结果到一个map<value, count>里面，然后用一个大小为k的最小堆处理，每次从map中读取元素时根据count值判断是否要替换掉堆顶。由于STL的map是用红黑树实现的，所以时间复杂度就保证了，要再快的话就直接用hash函数吧，可以到O(n)。维护k大小的堆，算最差情况，建堆复杂度是O(k)，维护复杂度是O(logk), 要维护M-k次，其中0<M<=N是map中的元素个数.
代码

struct CountVal{
    int val;
    int count;
    CountVal(int _val, int _count){
        val = _val;
        count = _count;
    }
};
bool compareObj(const CountVal& x, const CountVal& y){
    return x.count > y.count;
}
class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        map<int, int> countMap;
        vector<CountVal> counts;
        for(int i=0; i<nums.size(); i++){
            if(countMap.count(nums[i]) == 0)
                countMap[nums[i]] = 1;
            else
                countMap[nums[i]]++;
        }
        int count = 0;
        for(auto it=countMap.begin(); it!=countMap.end(); it++){
            CountVal item(it->first, it->second);
            if(count < k){
                counts.push_back(item);
                if(count == k-1){
                    make_heap(counts.begin(), counts.end(), compareObj);
                }
            } else {
                if (it->second > counts.front().count) {
                    pop_heap(counts.begin(), counts.end(), compareObj);
                    counts.pop_back();
                    counts.push_back(item);
                    push_heap(counts.begin(), counts.end(), compareObj);
                }
            }
            count++;
        }
        vector<int> results;
        for(auto it=counts.begin(); it!=counts.end(); it++){
            results.push_back(it->val);
        }
        return results;
    }
};