Top K Frequent Elements Using Heap in DSA Go

Mental Model

Find the most common items by counting them, then keep only the top few using a special list that quickly removes the least common.

Analogy: Imagine you have a basket of fruits and want to keep only the top 2 most frequent fruits. You count each fruit, then keep adding fruits to a small box that always throws out the least common fruit when full.

Input array: [1,1,1,2,2,3]
Frequency map: {1:3, 2:2, 3:1}
Min-heap (size k=2):
  [2:2]
  [1:3]
Heap keeps smallest frequency on top to remove when new bigger frequency comes.

Dry Run Walkthrough

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

Goal: Find the 2 most frequent elements in the list

Step 1: Count frequency of each number

Frequency map: {1:3, 2:2, 3:1}

Why: We need to know how often each number appears to find the most frequent

Step 2: Add (1,3) to min-heap

Heap: [(1,3)]

Why: Start building heap with first frequency

Step 3: Add (2,2) to min-heap

Heap: [(2,2), (1,3)]

Why: Add second frequency; heap keeps smallest frequency on top

Step 4: Add (3,1) to min-heap; since heap size > k, remove smallest frequency

Heap before removal: [(3,1), (1,3), (2,2)]
Heap after removal: [(2,2), (1,3)]

Why: Keep heap size at k=2 by removing least frequent element

Step 5: Extract elements from heap as result

Result: [1, 2]

Why: These are the top 2 frequent elements

Result:

Heap final state: [(2,2), (1,3)]
Top 2 frequent elements: [1, 2]

Annotated Code

DSA Go

package main

import (
	"container/heap"
	"fmt"
)

// Item holds number and its frequency
type Item struct {
	value    int
	priority int
}

// MinHeap implements heap.Interface for Items based on priority
type MinHeap []Item

func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i].priority < h[j].priority }
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.(Item))
}

func (h *MinHeap) Pop() interface{} {
	n := len(*h)
	item := (*h)[n-1]
	*h = (*h)[:n-1]
	return item
}

func topKFrequent(nums []int, k int) []int {
	freqMap := make(map[int]int)
	for _, num := range nums {
		freqMap[num]++
	}

	h := &MinHeap{}
	heap.Init(h)

	for val, freq := range freqMap {
		heap.Push(h, Item{value: val, priority: freq})
		if h.Len() > k {
			heap.Pop(h) // remove least frequent
		}
	}

	result := make([]int, 0, k)
	for h.Len() > 0 {
		item := heap.Pop(h).(Item)
		result = append(result, item.value)
	}

	// reverse result because heap pops smallest first
	for i, j := 0, len(result)-1; i < j; i, j = i+1, j-1 {
		result[i], result[j] = result[j], result[i]
	}

	return result
}

func main() {
	nums := []int{1, 1, 1, 2, 2, 3}
	k := 2
	res := topKFrequent(nums, k)
	fmt.Println(res)
}

freqMap := make(map[int]int) for _, num := range nums { freqMap[num]++ }

Count frequency of each number in input

h := &MinHeap{} heap.Init(h)

Initialize empty min-heap to store top k frequent elements

for val, freq := range freqMap { heap.Push(h, Item{value: val, priority: freq}) if h.Len() > k { heap.Pop(h) // remove least frequent } }

Add each frequency to heap; keep heap size at k by removing smallest frequency

for h.Len() > 0 { item := heap.Pop(h).(Item) result = append(result, item.value) }

Extract elements from heap to build result list

for i, j := 0, len(result)-1; i < j; i, j = i+1, j-1 { result[i], result[j] = result[j], result[i] }

Reverse result to have most frequent elements first

Output

[1 2]

Complexity Analysis

Time: O(n log k) because we count frequencies in O(n) and maintain a heap of size k with O(log k) operations for each unique element

Space: O(n) for frequency map and heap storage in worst case when all elements are unique

vs Alternative: Compared to sorting all frequencies O(n log n), using a heap reduces time to O(n log k) which is faster when k is much smaller than n

Edge Cases

Empty input array

Returns empty list since no elements to process

DSA Go

freqMap := make(map[int]int)
for _, num := range nums {
	freqMap[num]++
}

k equals number of unique elements

Returns all unique elements sorted by frequency

DSA Go

if h.Len() > k {
	heap.Pop(h) // remove least frequent
}

All elements have same frequency

Returns any k elements since frequencies tie

DSA Go

heap.Push and heap.Pop maintain heap size and order

When to Use This Pattern

When asked to find the top k frequent items efficiently, use a min-heap to keep track of the k largest frequencies without sorting all elements.

Common Mistakes

Mistake: Using a max-heap and popping k times instead of maintaining a min-heap of size k
Fix: Use a min-heap to keep only k elements, popping the smallest frequency when size exceeds k

Mistake: Not reversing the result after popping from min-heap, resulting in least frequent first
Fix: Reverse the result list after extracting from heap to get most frequent first

Summary

Find the k most frequent elements by counting and using a min-heap to keep top k frequencies.

Use when you need the most common items efficiently without sorting all frequencies.

Keep a small heap of size k that always removes the least frequent to maintain top k.