Top K Frequent Elements - Leetcode Solution

Problem Link

Step-by-Step Thought Process

Brute Force

Understand the problem: Find the k most frequent elements in an array of integers.
Use a Counter to count the frequency of each number in nums.
Initialize an empty min-heap heap to store at most k (frequency, number) pairs.
Iterate through each key (number) and val (frequency) in the Counter.
If the heap has fewer than k elements, push (val, key) to the heap.
Otherwise, push (val, key) to the heap and pop the smallest element to maintain k elements.
Extract the second element (number) from each tuple in the heap and return them as the result.

Code Solution (Brute Force)

from collections import Counter
import heapq
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        counter = Counter(nums)
        heap = []

        for key, val in counter.items():
            if len(heap) < k:
                heapq.heappush(heap, (val, key))
            else:
                heapq.heappushpop(heap, (val, key))
        
        return [h[1] for h in heap]
    # Time: O(n log k), Space: O(k)

#include <vector>
#include <unordered_map>
#include <queue>
#include <algorithm>
using namespace std;

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> countMap;
        for (int num : nums) {
            countMap[num]++;
        }
        
        auto cmp = [](const pair<int, int>& a, const pair<int, int>& b) {
            return a.second > b.second;
        };
        
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> minHeap(cmp);
        
        for (auto& entry : countMap) {
            if (minHeap.size() < k) {
                minHeap.push(entry);
            } else {
                minHeap.push(entry);
                minHeap.pop();
            }
        }
        
        vector<int> topK;
        while (!minHeap.empty()) {
            topK.push_back(minHeap.top().first);
            minHeap.pop();
        }
        
        reverse(topK.begin(), topK.end());
        return topK;
    }
};

import java.util.*;

public class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> countMap = new HashMap<>();
        for (int num : nums) {
            countMap.put(num, countMap.getOrDefault(num, 0) + 1);
        }
        
        PriorityQueue<Map.Entry<Integer, Integer>> heap = new PriorityQueue<>(
            (a, b) -> Integer.compare(a.getValue(), b.getValue())
        );
        
        for (Map.Entry<Integer, Integer> entry : countMap.entrySet()) {
            if (heap.size() < k) {
                heap.offer(entry);
            } else {
                heap.offer(entry);
                heap.poll();
            }
        }
        
        int[] topK = new int[k];
        for (int i = k - 1; i >= 0; i--) {
            topK[i] = heap.poll().getKey();
        }
        
        return topK;
    }
}

var topKFrequent = function(nums, k) {
    const countMap = new Map();
    for (const num of nums) {
        countMap.set(num, (countMap.get(num) || 0) + 1);
    }

    const heap = new MinHeap();

    for (const [num, freq] of countMap) {
        if (heap.size() < k) {
            heap.add([freq, num]);
        } else {
            heap.add([freq, num]);
            heap.poll();
        }
    }

    return heap.toArray().map(pair => pair[1]);
};

// Helper class for MinHeap
class MinHeap {
    constructor() {
        this.heap = [];
    }

    size() {
        return this.heap.length;
    }

    add(pair) {
        this.heap.push(pair);
        this.bubbleUp();
    }

    poll() {
        const min = this.heap[0];
        const end = this.heap.pop();
        if (this.size() > 0) {
            this.heap[0] = end;
            this.bubbleDown();
        }
        return min;
    }

    bubbleUp() {
        let idx = this.size() - 1;
        const element = this.heap[idx];
        while (idx > 0) {
            let parentIdx = Math.floor((idx - 1) / 2);
            let parent = this.heap[parentIdx];
            if (element[0] >= parent[0]) break;
            this.heap[idx] = parent;
            idx = parentIdx;
        }
        this.heap[idx] = element;
    }

    bubbleDown() {
        let idx = 0;
        const length = this.size();
        const element = this.heap[0];
        while (true) {
            let leftChildIdx = 2 * idx + 1;
            let rightChildIdx = 2 * idx + 2;
            let swap = null;

            if (leftChildIdx < length) {
                let leftChild = this.heap[leftChildIdx];
                if (leftChild[0] < element[0]) {
                    swap = leftChildIdx;
                }
            }

            if (rightChildIdx < length) {
                let rightChild = this.heap[rightChildIdx];
                if ((swap === null && rightChild[0] < element[0]) ||
                    (swap !== null && rightChild[0] < this.heap[swap][0])) {
                    swap = rightChildIdx;
                }
            }

            if (swap === null) break;
            this.heap[idx] = this.heap[swap];
            idx = swap;
        }
        this.heap[idx] = element;
    }

    toArray() {
        return this.heap;
    }
}

Why the Brute-Force Solution is Inefficient

The brute-force solution uses a heap to maintain the k most frequent elements, leading to:

Time Complexity: O(n log k), where n is the length of the array, due to heap operations for each of the n elements (or fewer distinct elements).
Space Complexity: O(k), for storing the heap of k elements.
Performance Issue: While efficient for small k, the heap operations introduce a logarithmic factor that can be avoided by leveraging bucket sort, which is more suitable given the frequency range is bounded by n.

Optimal Solution

The optimal solution uses bucket sort to achieve O(n) time complexity:

Get the length n of the input array nums.
Use a Counter to count the frequency of each number in nums.
Initialize a buckets array of size n+1, where buckets[i] stores a list of numbers with frequency i.
Iterate through the Counter, placing each number in buckets[freq]; if buckets[freq] is empty, initialize it as a list with the number, otherwise append the number.
Initialize an empty list ret to store the result.
Iterate through buckets from index n down to 0; if buckets[i] is not empty, extend ret with buckets[i].
Stop when ret contains k elements and return ret.

Code Solution (Optimal)

from collections import Counter
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        n = len(nums)
        counter = Counter(nums)
        buckets = [0] * (n + 1)

        for num, freq in counter.items():
            if buckets[freq] == 0:
                buckets[freq] = [num]
            else:
                buckets[freq].append(num)
        
        ret = []
        for i in range(n, -1, -1):
            if buckets[i] != 0:
                ret.extend(buckets[i])
            if len(ret) == k:
                break
        
        return ret

# Time Complexity: O(n)
# Space Complexity: O(n)

#include <vector>
#include <unordered_map>
#include <unordered_set>

using namespace std;

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> count;
        for (int num : nums) {
            count[num]++;
        }

        vector<vector<int>> buckets(nums.size() + 1);
        for (auto& entry : count) {
            buckets[entry.second].push_back(entry.first);
        }

        vector<int> result;
        for (int i = buckets.size() - 1; i >= 0; --i) {
            if (!buckets[i].empty()) {
                result.insert(result.end(), buckets[i].begin(), buckets[i].end());
                if (result.size() >= k) {
                    result.resize(k);
                    break;
                }
            }
        }

        return result;
    }
};

import java.util.*;

public class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> count = new HashMap<>();
        for (int num : nums) {
            count.put(num, count.getOrDefault(num, 0) + 1);
        }

        List<Integer>[] buckets = new ArrayList[nums.length + 1];
        for (int i = 0; i < buckets.length; i++) {
            buckets[i] = new ArrayList<>();
        }

        for (Map.Entry<Integer, Integer> entry : count.entrySet()) {
            int num = entry.getKey();
            int freq = entry.getValue();
            buckets[freq].add(num);
        }

        List<Integer> result = new ArrayList<>();
        for (int i = buckets.length - 1; i >= 0; --i) {
            if (!buckets[i].isEmpty()) {
                result.addAll(buckets[i]);
                if (result.size() >= k) {
                    break;
                }
            }
        }

        // Convert result to int[]
        int[] resultArray = new int[k];
        for (int i = 0; i < k; i++) {
            resultArray[i] = result.get(i);
        }
        return resultArray;
    }
}

var topKFrequent = function(nums, k) {
    const count = new Map();
    for (const num of nums) {
        count.set(num, (count.get(num) || 0) + 1);
    }

    const buckets = Array.from({ length: nums.length + 1 }, () => []);
    for (const [num, freq] of count.entries()) {
        buckets[freq].push(num);
    }

    const result = [];
    for (let i = buckets.length - 1; i >= 0; i--) {
        if (buckets[i].length > 0) {
            result.push(...buckets[i]);
            if (result.length >= k) {
                return result.slice(0, k);
            }
        }
    }
};

Detailed Explanation

Understanding the Problem: Top K Frequent Elements

In this problem, we're given an unsorted array of integers and asked to return the k elements that appear most frequently. This problem is a classic example of frequency analysis and appears frequently in scenarios such as log data analysis, recommendation systems, and natural language processing where you need to extract high-occurrence entities.

Initial Approach: Min Heap

A common brute-force approach involves counting the frequency of each element and using a min heap to keep track of the top k most frequent numbers. We build a frequency map using a hash table (or Python’s Counter) and then push each (frequency, number) pair into the heap. When the heap size exceeds k, we remove the smallest frequency element. This ensures the heap always contains the k most frequent elements seen so far.

Once all elements are processed, the numbers in the heap are our result. The heap allows for efficient retrieval of the most frequent items in O(n log k) time.

Optimized Strategy: Bucket Sort

While the heap-based approach is efficient, we can do better. The maximum frequency any number can have is equal to the length of the input array. So instead of sorting or using a heap, we can use a bucket sort idea. We create an array of n + 1 buckets where index i represents all numbers that appear exactly i times.

After populating the buckets, we iterate from the highest possible frequency down to 1, collecting numbers until we have gathered k of them. This approach skips sorting altogether and guarantees a linear time solution in O(n).

Time and Space Complexity

The heap solution has a time complexity of O(n log k), where n is the number of unique elements and k is the number of top elements to retrieve. It uses O(k) space for the heap and O(n) for the frequency map.

The bucket sort solution improves the time complexity to O(n) in the average case by eliminating the log factor introduced by the heap. The space complexity remains O(n) due to the use of the frequency map and buckets.

Conclusion

The “Top K Frequent Elements” problem showcases the trade-offs between general-purpose data structures like heaps and specialized techniques like bucket sort. The heap approach is versatile and works well in many scenarios, but if the data’s frequency bounds are known or small, bucket sort offers superior performance. Understanding both strategies enhances your problem-solving toolkit for real-world frequency analysis problems.

Get Personalized Support at AlgoMap Bootcamp 💡