4.3 Bucket Sort

347. Top K Frequent Elements

Problem Description

Given an array, find the top $k$ most frequent elements.

Input and Output Example

Input is an array and a target value $k$. Output is an array of length $k$.

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

In this example, the two most frequent numbers are 1 and 2.

Solution Explanation

As the name suggests, bucket sort involves creating a bucket for each value, where the bucket records the frequency of that value (or other attributes), and then sorting the buckets. For this example, we first use bucket sort to obtain four buckets [1,2,3,4], with values [4,2,1,1], representing the frequency of each number.

Next, we sort the buckets by frequency, and the top $k$ buckets are the top $k$ most frequent elements. We can use any sorting algorithm here or even perform another bucket sort, placing each old bucket into new buckets based on frequency. For this example, since the maximum frequency is 4, we create four new buckets [1,2,3,4], which contain the old buckets as follows: [[3,4],[2],[],[1]]. Finally, we iterate from the end to the beginning until we find $k$ old buckets.

We can use unordered_map in C++ or dict in Python to implement a hash table.

C++

vector<int> topKFrequent(vector<int>& nums, int k) {
    unordered_map<int, int> counts;
    for (int num : nums) {
        ++counts[num];
    }
    unordered_map<int, vector<int>> buckets;
    for (auto [num, count] : counts) {
        buckets[count].push_back(num);
    }
    vector<int> top_k;
    for (int count = nums.size(); count >= 0; --count) {
        if (buckets.contains(count)) {
            for (int num : buckets[count]) {
                top_k.push_back(num);
                if (top_k.size() == k) {
                    return top_k;
                }
            }
        }
    }
    return top_k;
}

Python

def topKFrequent(nums: List[int], k: int) -> List[int]:
    counts = Counter(nums)
    buckets = dict()
    for num, count in counts.items():
        if count in buckets:
            buckets[count].append(num)
        else:
            buckets[count] = [num]
    top_k = []
    for count in range(len(nums), 0, -1):
        if count in buckets:
            top_k += buckets[count]
            if len(top_k) >= k:
                return top_k[:k]
    return top_k[:k]