Top K Frequent Elements Using Heap in DSA Python

Mental Model

We find the elements that appear most often by counting them, then keep only the top K using a small container that quickly removes the least frequent.

Analogy: Imagine you have a bag of colored balls and want to keep only the K colors you see most often. You count how many balls of each color you have, then keep adding colors to a small box that always throws out the color with the fewest balls when full.

Input array: [1,1,1,2,2,3]
Frequency map: {1:3, 2:2, 3:1}
Min-heap (size K=2):
  (freq=2, val=2)
  (freq=3, val=1)
Top K frequent elements: [1, 2]

Dry Run Walkthrough

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

Goal: Find the 2 elements that appear most frequently

Step 1: Count frequency of each element

{1:3, 2:2, 3:1}

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

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

Heap: [(3,1)]

Why: Start building heap with first element frequency

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

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

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

Step 4: Add (freq=1, val=3) to min-heap; heap size exceeds k=2, remove smallest

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

Why: Keep heap size at k by removing element with lowest frequency

Step 5: Extract elements from heap as top k frequent

Top k elements: [2, 1]

Why: Heap now contains the k most frequent elements

Result:

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

Annotated Code

DSA Python

import heapq
from typing import List

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        freq_map = {}
        for num in nums:
            freq_map[num] = freq_map.get(num, 0) + 1

        min_heap = []
        for num, freq in freq_map.items():
            heapq.heappush(min_heap, (freq, num))  # push frequency and number
            if len(min_heap) > k:
                heapq.heappop(min_heap)  # remove smallest frequency

        result = []
        while min_heap:
            freq, num = heapq.heappop(min_heap)
            result.append(num)

        return result[::-1]  # reverse to get highest frequency first

if __name__ == "__main__":
    nums = [1,1,1,2,2,3]
    k = 2
    sol = Solution()
    top_k = sol.topKFrequent(nums, k)
    print(" -> ".join(map(str, top_k)) + " -> null")

freq_map[num] = freq_map.get(num, 0) + 1

count frequency of each element

heapq.heappush(min_heap, (freq, num)) # push frequency and number

add element frequency to min-heap

if len(min_heap) > k: heapq.heappop(min_heap) # remove smallest frequency

keep heap size at k by removing least frequent element

freq, num = heapq.heappop(min_heap) result.append(num)

extract elements from heap to build result

return result[::-1] # reverse to get highest frequency first

reverse list to have most frequent elements first

Output

1 -> 2 -> null

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) per insertion

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

vs Alternative: Using sorting all frequencies would be O(n log n), heap reduces this to O(n log k) which is faster when k is much smaller than n

Edge Cases

empty array

returns empty list since no elements to count

DSA Python

freq_map = {}

k equals number of unique elements

returns all unique elements sorted by frequency

DSA Python

if len(min_heap) > k:
    heapq.heappop(min_heap)

all elements have same frequency

returns any k elements since frequencies tie

DSA Python

heapq.heappush(min_heap, (freq, num))

When to Use This Pattern

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

Common Mistakes

Mistake: Using a max-heap and popping all elements to get top K, which is less efficient
Fix: Use a min-heap of size K to keep only top K elements, popping smallest frequency when heap grows

Mistake: Not reversing the result list after popping from min-heap, resulting in ascending frequency order
Fix: Reverse the result list before returning to have highest frequency first

Summary

Find the K elements that appear most frequently using a frequency map and a min-heap.

Use when you need the top K frequent elements efficiently without sorting all data.

Keep a min-heap of size K to quickly discard less frequent elements and keep only the top K.

Practice

1. What is the main purpose of using a heap in the Top K Frequent Elements problem?

easy

A. To efficiently keep track of the top k elements by frequency

B. To sort all elements in ascending order

C. To store elements without any order

D. To remove duplicates from the list

Solution

1. Step 1: Understand the problem goal

   The problem asks to find the k elements that appear most frequently in a list.

2. Step 2: Role of heap in this context

   A heap helps efficiently keep track of the top k frequent elements by maintaining a small structure of size k, allowing quick access to the smallest frequency among them.

3. Final Answer:

   To efficiently keep track of the top k elements by frequency -> Option A

4. Quick Check:

   Heap = Efficient top k tracking [OK]

Hint: Heap keeps smallest frequency on top for quick replacement [OK]

Common Mistakes:

• Thinking heap sorts all elements fully
• Confusing heap with simple list or set
• Assuming heap removes duplicates

2. Which of the following Python code snippets correctly creates a min-heap from a list freq_list using the heapq module?

easy

A. import heapq
heapq.makeheap(freq_list)

B. import heapq
heapq.heapify(freq_list)

C. import heapq
heapq.heap(freq_list)

D. import heapq
heapq.create(freq_list)

Solution

1. Step 1: Recall heapq module functions

   The heapq module provides heapify() to convert a list into a heap in-place.

2. Step 2: Identify correct function usage

   Only heapq.heapify(freq_list) is valid syntax to create a min-heap from a list.

3. Final Answer:

   import heapq heapq.heapify(freq_list) -> Option B

4. Quick Check:

   heapify() creates heap [OK]

Hint: Use heapq.heapify() to build a heap from list [OK]

Common Mistakes:

• Using non-existent heapq functions
• Confusing heapify with sorting
• Missing import statement

3. What will be the output of the following code?

import heapq
nums = [1,1,1,2,2,3]
count = {}
for n in nums:
    count[n] = count.get(n, 0) + 1
heap = []
for num, freq in count.items():
    heapq.heappush(heap, (freq, num))
result = []
for _ in range(2):
    result.append(heapq.heappop(heap)[1])
print(result)

medium

A. [1, 2]

B. [2, 3]

C. [1, 3]

D. [3, 2]

Solution

1. Step 1: Count frequencies of elements

   nums = [1,1,1,2,2,3] gives count = {1:3, 2:2, 3:1}.

2. Step 2: Push (freq, num) into min-heap and pop twice

   Heap contains (3,1), (2,2), (1,3). Popping twice removes smallest freq first: (1,3), then (2,2). So result = [3, 2].

3. Final Answer:

   [3, 2] -> Option D

4. Quick Check:

   Lowest freq elements popped first = [3, 2] [OK]

Hint: Min-heap pops smallest frequency first [OK]

Common Mistakes:

• Assuming heap pops highest frequency first
• Mixing up element and frequency order
• Not counting frequencies correctly

4. Identify the error in the following code snippet that tries to find the top 2 frequent elements:

import heapq
nums = [4,4,4,5,5,6]
count = {}
for n in nums:
    count[n] = count.get(n, 0) + 1
heap = []
for num, freq in count.items():
    heapq.heappush(heap, (num, freq))
result = []
for _ in range(2):
    result.append(heapq.heappop(heap)[1])
print(result)

medium

A. The range in the last loop should be 3 instead of 2

B. The count dictionary is not updated correctly

C. The tuple pushed to heap should be (freq, num), not (num, freq)

D. heapq.heappush is used incorrectly; should use heapq.push

Solution

1. Step 1: Check tuple order pushed to heap

   The heap is a min-heap by first tuple element. Pushing (num, freq) orders by num, not frequency.

2. Step 2: Correct tuple order for frequency sorting

   To get top frequent elements, push (freq, num) so heap orders by frequency.

3. Final Answer:

   The tuple pushed to heap should be (freq, num), not (num, freq) -> Option C

4. Quick Check:

   Heap order depends on first tuple element [OK]

Hint: Push (frequency, element) tuples to heap [OK]

Common Mistakes:

• Swapping frequency and element in tuple
• Using wrong heapq function name
• Incorrect loop range for popping

5. Given the list nums = [2,3,2,4,3,3,5,5,5,5] and k = 2, which of the following code snippets correctly returns the top k frequent elements using a heap?

hard

A. import heapq
count = {}
for n in nums:
count[n] = count.get(n, 0) + 1
heap = []
for num, freq in count.items():
if len(heap) < k:
heapq.heappush(heap, (freq, num))
else:
if freq > heap[0][0]:
heapq.heapreplace(heap, (freq, num))
result = [num for freq, num in heap]
print(result)

B. import heapq
count = {}
for n in nums:
count[n] = count.get(n, 0) + 1
heap = []
for num, freq in count.items():
heapq.heappush(heap, (num, freq))
result = [heapq.heappop(heap)[1] for _ in range(k)]
print(result)

C. import heapq
count = {}
for n in nums:
count[n] = count.get(n, 0) + 1
heap = []
for num, freq in count.items():
heapq.heappush(heap, (-freq, num))
result = [heapq.heappop(heap)[0] for _ in range(k)]
print(result)

D. import heapq
count = {}
for n in nums:
count[n] = count.get(n, 0) + 1
heap = []
for num, freq in count.items():
if freq < k:
heapq.heappush(heap, (freq, num))
result = [num for freq, num in heap]
print(result)

Solution

1. Step 1: Understand the problem and approach

   We want top k frequent elements using a min-heap of size k to keep track of highest frequencies.

2. Step 2: Analyze each option

   import heapq count = {} for n in nums: count[n] = count.get(n, 0) + 1 heap = [] for num, freq in count.items(): if len(heap) < k: heapq.heappush(heap, (freq, num)) else: if freq > heap[0][0]: heapq.heapreplace(heap, (freq, num)) result = [num for freq, num in heap] print(result) correctly pushes (freq, num) and maintains heap size k, replacing smallest freq when a bigger freq is found. import heapq count = {} for n in nums: count[n] = count.get(n, 0) + 1 heap = [] for num, freq in count.items(): heapq.heappush(heap, (num, freq)) result = [heapq.heappop(heap)[1] for _ in range(k)] print(result) pushes (num, freq) which is wrong order. import heapq count = {} for n in nums: count[n] = count.get(n, 0) + 1 heap = [] for num, freq in count.items(): heapq.heappush(heap, (-freq, num)) result = [heapq.heappop(heap)[1] for _ in range(k)] print(result) uses negative freq but does not maintain heap size k properly. import heapq count = {} for n in nums: count[n] = count.get(n, 0) + 1 heap = [] for num, freq in count.items(): if freq < k: heapq.heappush(heap, (freq, num)) result = [num for freq, num in heap] print(result) pushes only if freq < k which is incorrect logic.

3. Final Answer:

   Option A code correctly returns top k frequent elements -> Option A

4. Quick Check:

   Min-heap size k with freq as key [OK]

Hint: Maintain min-heap size k with (freq, num) tuples [OK]

Common Mistakes:

• Pushing wrong tuple order to heap
• Not limiting heap size to k
• Using negative frequencies without proper logic