550. Top K Frequent Words II

map<string, int> words; class cmp { public: bool operator () (const string& a, const string& b) { int a_count = words[a]; int b_count = words[b]; if (a_count != b_count) return a_count > b_count; return a < b; } }; class TopK { private: set<string, cmp> q; int k; public: TopK(int k) { // initialize your data structure here this->k = k; } void add(string& word) { // Write your code here if (words.find(word) == words.end()) { words[word] = 1; } else { if (q.find(word) != q.end()) q.erase(word); words[word] += 1; } q.insert(word); if (q.size() > k) q.erase(--q.end()); } vector<string> topk() { // Write your code here vector<string> topk; //set<string, cmp>::iterator it = q.begin(); int i = 0; for (auto it : q) { topk.push_back(it); } return topk; } };