3. Top K Frequent Elements

Link to the problem

Description

Given a non-empty array of integers, return the k most frequent elements.

Note:

• You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
• Your algorithm's time complexity must be better than O(n log n), where n is the array's size.

Example

Input: [1,1,1,2,2,3], k = 2, Output: [1,2].
1 is most frequent, 2 is the second frequent element.

Idea

First construct a counter of the elements, then put the elements in an array indexed by the count. Finally, read the array from right to left.

Solution

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        int n = nums.size();
        unordered_map<int, int> counter;
        vector<vector<int> > counts(n + 1);
        vector<int> rtn;
        for (auto it = nums.begin(); it != nums.end(); it++) {
            int key = *it;
            if (counter.find(key) == counter.end()) counter[key] = 1;
            else counter[key]++;
        }
        for (auto it = counter.begin(); it != counter.end(); it++) {
            int element = it->first, count = it->second;
            counts[count].push_back(element);
        }
        int index = n;
        while (rtn.size() < k) {
            for (auto it = counts[index].begin(); it != counts[index].end(); it++) {
                if (rtn.size() == k) return rtn;
                rtn.push_back(*it);
            }
            index--;
        }
        return rtn;
    }
};

Performance

20 / 20 test cases passed.
Runtime: 19 ms