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]
For LeetCode 347 – Top K Frequent Elements. We will aim to achieve an O(n log k) time complexity, which is the most efficient for this problem.
Key steps to solve this problem:
- Count the frequency of each element in the array using a map.
- Use a min-heap (priority queue) to keep track of the top
k
frequent elements. The heap will store pairs of(element, frequency)
. - Extract the top k elements from the heap.
Code implementation:
package main
import (
"container/heap"
"fmt"
)
// Define a type for the min-heap
type Element struct {
num int
freq int
}
type MinHeap []Element
// Implementing heap.Interface methods
func (h MinHeap) Len() int { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i].freq < h[j].freq }
func (h MinHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) {
*h = append(*h, x.(Element))
}
func (h *MinHeap) Pop() interface{} {
old := *h
n := len(old)
item := old[n-1]
*h = old[0 : n-1]
return item
}
func topKFrequent(nums []int, k int) []int {
// Step 1: Build frequency map
freqMap := make(map[int]int)
for _, num := range nums {
freqMap[num]++
}
// Step 2: Use a min-heap to store the top k frequent elements
minHeap := &MinHeap{}
heap.Init(minHeap)
// Step 3: Iterate through the frequency map and maintain the heap
for num, freq := range freqMap {
heap.Push(minHeap, Element{num, freq})
if minHeap.Len() > k {
heap.Pop(minHeap) // Remove the element with the least frequency
}
}
// Step 4: Extract the top k frequent elements from the heap
result := make([]int, k)
for i := 0; i < k; i++ {
result[i] = heap.Pop(minHeap).(Element).num
}
return result
}
func main() {
// Test the function with an example
nums := []int{1,1,1,2,2,3}
k := 2
result := topKFrequent(nums, k)
fmt.Println(result) // Output: [1, 2]
}
Explanation:
- Frequency Map:
- We first create a frequency map
freqMap
to count how often each number appears in the input arraynums
.
- We first create a frequency map
- Min-Heap:
- We then use a min-heap to keep track of the top
k
frequent elements. A heap allows us to efficiently access and remove the least frequent element in O(log k) time.
- We then use a min-heap to keep track of the top
- Heap Operations:
- We push each element from the frequency map into the heap. If the size of the heap exceeds
k
, we pop the smallest element (which is the least frequent element so far).
- We push each element from the frequency map into the heap. If the size of the heap exceeds
- Result Extraction:
- After maintaining the heap with the top
k
frequent elements, we extract the elements from the heap into the result array.
- After maintaining the heap with the top
Time Complexity:
- Building the frequency map: O(n), where
n
is the number of elements innums
. - Heap operations: For each unique element in the frequency map, we perform heap operations (push and pop). In the worst case, the heap size is
k
, so each operation takes O(log k) time. There are at mostn
unique elements, so the heap operations take O(n log k) time. - Final extraction: O(k) to extract
k
elements from the heap.
Thus, the overall time complexity is O(n log k).
Space Complexity:
- O(n) for the frequency map.
- O(k) for the heap. Hence, the total space complexity is O(n + k).