Top K Frequent Elements Using Heap in Java

When faced with the problem of finding the top K frequent elements from a collection of items, the brute-force approach of counting frequencies followed by sorting can often be inefficient, especially as the dataset grows. Instead, leveraging a heap can significantly enhance performance. In this post, we'll delve into how to use a min-heap to efficiently solve the problem in Java.

Problem Statement

Given an integer array nums and an integer k, you are to return the k most frequent elements in the array.

Example

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

In this example, the element 1 appears three times, whereas 2 appears twice. Thus, the top two frequent elements are 1 and 2.

Using a HashMap and a Min-Heap

Step-by-Step Approach

1. Count Frequencies: First, we need a way to count how often each element appears in the array. A HashMap is ideal for this.

2. Use a Min-Heap: By maintaining a heap of size k, we can ensure that we only keep the top K frequent elements in our memory.

3. Extract Results: Finally, we can extract the elements from the heap to return our result.

Implementing the Solution

Let's implement these steps in Java:

import java.util.*;

public class TopKFrequentElements {
    
    public List<Integer> topKFrequent(int[] nums, int k) {
        // Step 1: Count frequencies using HashMap
        Map<Integer, Integer> frequencyMap = new HashMap<>();
        for (int num : nums) {
            frequencyMap.put(num, frequencyMap.getOrDefault(num, 0) + 1);
        }
        
        // Step 2: Create a min-heap based on frequency
        PriorityQueue<Map.Entry<Integer, Integer>> minHeap =
            new PriorityQueue<>(Comparator.comparingInt(Map.Entry::getValue));
        
        for (Map.Entry<Integer, Integer> entry : frequencyMap.entrySet()) {
            minHeap.offer(entry);
            // If the size of the heap is greater than k, remove the smallest frequency
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }
        
        // Step 3: Extract the elements from the min-heap into a result list
        List<Integer> result = new ArrayList<>();
        while (!minHeap.isEmpty()) {
            result.add(minHeap.poll().getKey());
        }
        
        // The result should be reversed to maintain the correct order
        Collections.reverse(result);
        return result;
    }
    
    public static void main(String[] args) {
        TopKFrequentElements solution = new TopKFrequentElements();
        int[] nums = {1, 1, 1, 2, 2, 3};
        int k = 2;
        System.out.println(solution.topKFrequent(nums, k)); // Output: [1, 2]
    }
}

Explanation of the Code

• Frequency Map: The frequencyMap counts occurrences of each integer using a single pass through the nums array.

• Min-Heap: A priority queue (min-heap) is initialized where we store the frequencies. The comparator defined here sorts based on the count of occurrences.

• Maintaining Heap Size: We ensure the size of the heap doesn't exceed k. If it does, we remove the element with the least frequency, which allows us to retain only the top K frequent elements.

• Final List: We extract the keys from the heap (elements) into a result list and reverse it since the heap organized them in ascending order of frequency.

Complexity Analysis

• Time Complexity: Counting the frequencies takes O(n), where n is the length of nums. Each of the K elements is added to the heap takes O(log k) time, leading to an overall complexity of O(n log k).

• Space Complexity: The space used is O(n) for the frequency map and O(k) for the min-heap, resulting in an overall space complexity of O(n).

Conclusion

The use of a min-heap in conjunction with a frequency map allows us to efficiently determine the top K most frequent elements in an array. This method not only reduces complexity substantially but also showcases the versatility of heaps in solving various data structure and algorithm problems.

With practice, you can integrate this approach into your toolkit for handling similar array querying problems effectively!