topk.cpp

#include <vector>
#include <unordered_map>
#include <queue>

using namespace std;

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        vector<int> result;
        unordered_map<int, int> count;
        for (vector<int>::iterator itr = nums.begin(); itr != nums.end(); ++itr) {
            ++count[*itr];
        }
        
        priority_queue<pair<int, int>, vector<pair<int, int> >, std::greater<pair<int, int> > > topk;
        for (unordered_map<int, int>::iterator itr = count.begin(); itr != count.end(); ++itr) {
            topk.push(make_pair(itr->second, itr->first));
            if (topk.size() > k) {
                topk.pop();
            }
        }
        
        while (!topk.empty()) {
            result.push_back(topk.top().second);
            topk.pop();
        }
        reverse(result.begin(), result.end());
        
        return result;
    }
};