A std::priority_queue
uses a std::vector
as the default container (Reference this). For sorting on the basis of the first element in a std::vector<pair<int, int>>
, we need to define our own comparison function (Reference this). This is what I understand.
Now, the following code returns the k
most frequent elements in a non-empty array, in O(NlogK):
class Solution {
public:
vector<int> topKFrequent(vector<int>& nums, int k) {
if(nums.empty())
return vector<int>();
unordered_map< int, int > hashMap;
for(int i=0; i<nums.size(); i++)
hashMap[nums[i]]++;
priority_queue< pair< int, int >> pq;
vector< int > result;
unordered_map< int, int >::iterator it=hashMap.begin();
for(it=hashMap.begin(); it!=hashMap.end(); it++) {
//the first one is frequency and the second one is the value
pq.push(make_pair(it->second, it->first));
//the peculiar implementation below is because we the code to be O(NlogK)
if(pq.size()>(hashMap.size()-k)) {
result.push_back(pq.top().second);
pq.pop();
}
}
return result;
}
};
This code works correctly and gets accepted by the judge - but how? The std::priority_queue
, using a std::vector<pair<int, int>>
as its underlying container must contain a custom comparison function so that it sorts correctly. So, how does it work?