347. Top K Frequent Elements

Description

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

 

Example 1:

Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]

Example 2:

Input: nums = [1], k = 1
Output: [1]

 

Constraints:

• 1 <= nums.length <= 105
• -104 <= nums[i] <= 104
• k is in the range [1, the number of unique elements in the array].
• It is guaranteed that the answer is unique.

 

Follow up: Your algorithm's time complexity must be better than O(n log n), where n is the array's size.

Solution

top-k-frequent-elements.py

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        buckets = [[] for _ in range(len(nums)+1)]

        c = collections.Counter(nums).items()

        for num,freq in c:
            buckets[freq].append(num)

        res = []
        pointer = len(buckets) - 1

        while k > 0 and pointer >= 0:
            while k > 0 and len(buckets[pointer]) > 0:
                res.append(buckets[pointer].pop())
                k -= 1
            pointer -= 1

        return res

top-k-frequent-elements.cpp

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int,int> mp;

        for (int num: nums)
            ++mp[num];


        vector<vector<int>> buckets(nums.size() + 1); 
        for (auto v: mp)
            buckets[v.second].push_back(v.first);


        vector<int> res;
        for (int i = buckets.size() - 1; i >= 0 && res.size() < k; --i){
            for (int num : buckets[i]){
                res.push_back(num);

                if (res.size() == k)
                    break;
            }
        }

        return res;
    }
};