#include <vector>
#include <unordered_map>
 
// This is O(N) but there is still room for optimization.
 
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::vector<std::vector<int>> b(nums.size() + 1);
        for (const auto& [v, f]: v2f) b[f].push_back(v);
 
        // get the top k elements
        std::vector<int> result;
        result.reserve(k);
 
        bool breakEarly = false;
 
        for (int f = b.size() - 1; f > 0; f--) {
            for (const int& v: b[f]) {
                result.push_back(v);
 
                if (result.size() >= k) {
                    breakEarly = true;
                    break;
                }
            }
 
            if (breakEarly) break;
        }
 
        return result;
    }
};