347_top_k_frequent_elements_nlogn.cc

#include <vector>
#include <unordered_map>
#include <set>
#include <functional>
 
//  we basically do:
//  MAP-1 { v1 -> f1, ... } => MAP-2 { f1 -> [v1, v3, v6, ...], ... } ; SET { f1 > f3 > f2 > f5 > ... }
//  [x: x belongs to MAP-2[f], f belongs to SET, till we have k elements]
 
// This is almost an O(N) approach
 
class Solution {
public:
    std::vector<int> topKFrequent(std::vector<int>& nums, int k) {
        if (nums.size() <= 1) return nums;
 
        // kind of a bin sort but on frequencies
        std::unordered_map<int, int> v2f;
        for (const int& n: nums) v2f[n]++;
 
        std::unordered_map<int, std::vector<int>> f2v;
        std::vector<int> freqsTemp;
 
        for (const auto& [v, f]: v2f) {
            f2v[f].push_back(v);
            freqsTemp.push_back(f);
        }
        
        // This conversion to set is the O( N log N), rest of the code is O(N)
        std::set<int, std::greater<int>> freqs(freqsTemp.begin(), freqsTemp.end());
        std::vector<int> result;
        result.reserve(k);
 
        // get k most frequent elements
        bool breakEarly = false;
 
        for (const int& f: freqs) {
            for (const int& v: f2v[f]) {
                result.push_back(v);
 
                if (result.size() >= k) {
                    breakEarly = true;
                    break;
                }
            }
 
            if (breakEarly) break;
        }
 
        return result;
    }
};