LeetCode 347 – Top K Frequent Elements

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:

1. Count the frequency of each element in the array using a map.
2. Use a min-heap (priority queue) to keep track of the top k frequent elements. The heap will store pairs of (element, frequency).
3. 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:

1. Frequency Map:
   • We first create a frequency map freqMap to count how often each number appears in the input array nums.
2. 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.
3. 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).
4. Result Extraction:
   • After maintaining the heap with the top k frequent elements, we extract the elements from the heap into the result array.

Time Complexity:

• Building the frequency map: O(n), where n is the number of elements in nums.
• 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 most n 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).