M.T.Ramkrushna

Posted on Oct 3

Topic 11: Heaps & Priority Queues 🏔️

#ai #python #dsa

Heaps are extremely common in coding interviews because they let you efficiently handle problems involving:

Minimum / maximum values
Top K elements
Kth largest / smallest
Scheduling
Merging sorted data
Priority queues

In Python, the main tool is heapq.

1. What is a Heap?

A heap is a special complete binary tree.

There are two common types:

Min Heap

The smallest element is always at the root.

Max Heap

The largest element is always at the root.

Python's heapq directly provides a min heap, not a max heap.

2. Python `heapq`

import heapq

Create a heap:

heap = []

Add elements:

heapq.heappush(heap, 5)
heapq.heappush(heap, 2)
heapq.heappush(heap, 8)
heapq.heappush(heap, 1)

Now:

print(heap)

The smallest element is at:

heap[0]

Remove the smallest:

x = heapq.heappop(heap)

3. Basic Heap Operations

Operation	Complexity
Get minimum	`O(1)`
Insert	`O(log n)`
Remove minimum	`O(log n)`
Build heap	`O(n)`

Example:

import heapq

heap = []

heapq.heappush(heap, 5)
heapq.heappush(heap, 2)
heapq.heappush(heap, 8)
heapq.heappush(heap, 1)

print(heap[0])          # 1
print(heapq.heappop(heap))  # 1
print(heapq.heappop(heap))  # 2

4. Build a Heap from an Array

You can convert a list into a heap:

import heapq

nums = [5, 2, 8, 1, 3]

heapq.heapify(nums)

print(nums)

heapify() takes:

O(n)

This is an important interview fact.

5. Heap Is NOT a Sorted Array

This is a common beginner mistake.

After:

heapq.heapify(nums)

you cannot assume:

nums[0] <= nums[1] <= nums[2] ...

Only this is guaranteed:

nums[0] = smallest element

The remaining elements satisfy the heap property, not complete sorted order.

6. Max Heap in Python

Python doesn't have a direct max heap in heapq.

The standard trick is to use negative values.

import heapq

heap = []

heapq.heappush(heap, -10)
heapq.heappush(heap, -5)
heapq.heappush(heap, -20)

largest = -heapq.heappop(heap)

print(largest)

Output:

Why?

Python sees:

-20 < -10 < -5

So the smallest negative value corresponds to the largest original value.

7. Kth Largest Element

This is one of the most important heap interview problems.

Suppose:

nums = [3, 2, 1, 5, 6, 4]
k = 2

Answer:

because the sorted array is:

[1, 2, 3, 4, 5, 6]
             ↑

Efficient solution

Maintain a min heap of size k.

import heapq

def kth_largest(nums, k):
    heap = []

    for num in nums:
        heapq.heappush(heap, num)

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

    return heap[0]

Why does this work?

For:

k = 2

we keep only the two largest values seen so far.

At the end:

heap = [5, 6]

The smallest among these is:

which is the 2nd largest overall.

Complexity

Time:  O(n log k)
Space: O(k)

This is much better than sorting when k is small.

8. Top K Largest Elements

Suppose:

nums = [5, 1, 9, 3, 7, 2]
k = 3

We want:

[9, 7, 5]

Use a min heap of size k.

import heapq

def top_k_largest(nums, k):
    heap = []

    for num in nums:
        heapq.heappush(heap, num)

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

    return heap

The returned heap isn't necessarily sorted.

If sorted output is required:

return sorted(heap, reverse=True)

9. Kth Smallest Element

For the kth smallest, use a max heap of size k.

Since Python only has a min heap, store negatives.

import heapq

def kth_smallest(nums, k):
    heap = []

    for num in nums:
        heapq.heappush(heap, -num)

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

    return -heap[0]

Example:

nums = [7, 10, 4, 3, 20, 15]
k = 3

Sorted:

[3, 4, 7, 10, 15, 20]

Answer:

10. Priority Queue

A priority queue removes the item with the highest priority rather than simply the oldest item.

Python can implement this using a heap.

Example:

import heapq

tasks = []

heapq.heappush(tasks, (2, "Study"))
heapq.heappush(tasks, (1, "Interview"))
heapq.heappush(tasks, (3, "Exercise"))

priority, task = heapq.heappop(tasks)

print(task)

Output:

Interview

Because priority 1 comes first.

11. Heap with Tuples

Python compares tuples lexicographically.

heapq.heappush(heap, (priority, task))

It first compares:

priority

If priorities are equal, it compares:

task

Example:

heap = []

heapq.heappush(heap, (2, "C"))
heapq.heappush(heap, (1, "A"))
heapq.heappush(heap, (2, "B"))

print(heapq.heappop(heap))

Result:

(1, "A")

12. Custom Priority

Suppose you have:

students = [
    ("Alice", 90),
    ("Bob", 95),
    ("Charlie", 85)
]

If you want the highest score first, push negative scores:

import heapq

heap = []

for name, score in students:
    heapq.heappush(heap, (-score, name))

score, name = heapq.heappop(heap)

print(name, -score)

Output:

Bob 95

13. Merge K Sorted Lists

Very common interview problem.

Suppose:

List 1: 1 → 4 → 7
List 2: 2 → 5 → 8
List 3: 3 → 6 → 9

We want:

1 → 2 → 3 → 4 → 5 → 6 → 7 → 8 → 9

The heap stores the smallest current element from each list.

Conceptually:

heap
 ↓
1  2  3

Take 1, then add the next element from its list (4).

Now:

2  3  4

Continue.

Generic implementation

For arrays:

import heapq

def merge_k_sorted_lists(lists):
    heap = []

    for i, arr in enumerate(lists):
        if arr:
            heapq.heappush(heap, (arr[0], i, 0))

    result = []

    while heap:
        value, list_index, element_index = heapq.heappop(heap)
        result.append(value)

        next_index = element_index + 1

        if next_index < len(lists[list_index]):
            next_value = lists[list_index][next_index]

            heapq.heappush(
                heap,
                (next_value, list_index, next_index)
            )

    return result

Complexity

If there are N total elements and K lists:

Time:  O(N log K)
Space: O(K)

14. K Closest Points to Origin

Another famous heap problem.

Point:

(x, y)

Distance from origin:

x² + y²

We don't need the square root because comparing squared distances gives the same ordering.

Example:

points = [
    [1, 3],
    [-2, 2],
    [5, 8],
    [0, 1]
]

We want the k closest points.

A simple min-heap solution:

import heapq

def k_closest(points, k):
    heap = []

    for x, y in points:
        distance = x * x + y * y
        heapq.heappush(heap, (distance, x, y))

    result = []

    for _ in range(k):
        _, x, y = heapq.heappop(heap)
        result.append([x, y])

    return result

Complexity:

O(n log n)

There is also an O(n log k) max-heap approach, which is preferable when k is much smaller than n.

15. Find Median from Data Stream

This is a classic two heaps problem.

Maintain:

Max heap → smaller half
Min heap → larger half

Example:

Numbers:
1, 2, 3, 4

Conceptually:

max heap       min heap
[1, 2]   |   [3, 4]

Median:

(2 + 3) / 2 = 2.5

For:

1, 2, 3

we might have:

max heap       min heap
[1, 2]   |   [3]

Median:

This problem teaches an important pattern:

Use two heaps when you need to continuously maintain the middle of a changing dataset.

16. Heap vs Sorting

Suppose you need the largest k elements.

Sorting

nums.sort(reverse=True)
answer = nums[:k]

Complexity:

O(n log n)

Heap

O(n log k)

So if:

n = 1,000,000
k = 10

a heap can be much more appropriate.

17. When Should You Think "Heap"?

Look for phrases such as:

"K largest"

→ min heap of size K

"K smallest"

→ max heap of size K

"Repeatedly get minimum"

→ min heap

"Repeatedly get maximum"

→ max heap

"Priority"

→ priority queue / heap

"Merge K sorted"

→ min heap

"Median of a stream"

→ two heaps

"Next task with smallest cost/time"

→ priority queue

18. Important Heap Pattern

Memorize this:

Top K largest

heap = []

for x in nums:
    heapq.heappush(heap, x)

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

At the end:

heap[0]

is the kth largest.

Top K smallest

Use negative values:

heap = []

for x in nums:
    heapq.heappush(heap, -x)

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

Then:

-heap[0]

is the kth smallest.

19. Heap vs Stack vs Queue

Structure	Main behavior	Typical use
Stack	LIFO	DFS, parentheses
Queue	FIFO	BFS
Heap	Highest/lowest priority	Scheduling, Top K
Hash Map	Key → value	Fast lookup
Set	Unique values	Membership

This distinction is very important in interviews.

20. Common Mistakes

Mistake 1: Assuming heap is sorted

Wrong:

heapq.heapify(nums)
print(nums)

doesn't give a sorted list.

Mistake 2: Forgetting Python has a min heap

For a max heap, use:

-x

Mistake 3: Using a heap when sorting is simpler

If you need all elements sorted, simply sorting may be clearer.

Use heaps when you need to repeatedly access an extreme element or only need a subset such as Top K.

Mistake 4: Forgetting tuple tie-breaking

With:

heapq.heappush(heap, (priority, value))

Python compares priority first, then value if priorities are equal.

Interview Cheat Sheet 🧠

Minimum repeatedly
        ↓
    Min Heap

Maximum repeatedly
        ↓
    Max Heap

K largest
        ↓
Min Heap of size K

K smallest
        ↓
Max Heap of size K

Merge K sorted
        ↓
Min Heap

Median stream
        ↓
Two Heaps

Priority scheduling
        ↓
Priority Queue

Practice Questions

Easy

Implement a min heap using heapq
Kth Largest Element
Kth Smallest Element
Last Stone Weight
Top K Frequent Elements

Medium

K Closest Points to Origin
Merge K Sorted Lists
Task Scheduler
Find Median from Data Stream
Kth Largest Element in a Stream
Smallest Number in Infinite Set
Reorganize String

⭐ Must-do interview problems

If you're preparing specifically for coding interviews, make sure you can solve:

1. Kth Largest Element
2. Top K Frequent Elements
3. K Closest Points
4. Merge K Sorted Lists
5. Find Median from Data Stream

Next topic: Graphs — adjacency lists/matrices, DFS, BFS, connected components, cycle detection, and shortest paths.

1. What is a Heap?

Min Heap

Max Heap

2. Python heapq

3. Basic Heap Operations

4. Build a Heap from an Array

5. Heap Is NOT a Sorted Array

6. Max Heap in Python

7. Kth Largest Element

Efficient solution

Why does this work?

Complexity

8. Top K Largest Elements

9. Kth Smallest Element

10. Priority Queue

11. Heap with Tuples

12. Custom Priority

13. Merge K Sorted Lists

Generic implementation

Complexity

14. K Closest Points to Origin

15. Find Median from Data Stream

16. Heap vs Sorting

Sorting

Heap

17. When Should You Think "Heap"?

"K largest"

"K smallest"

"Repeatedly get minimum"

"Repeatedly get maximum"

"Priority"

"Merge K sorted"

"Median of a stream"

"Next task with smallest cost/time"

18. Important Heap Pattern

Top K largest

Top K smallest

19. Heap vs Stack vs Queue

20. Common Mistakes

Mistake 1: Assuming heap is sorted

Mistake 2: Forgetting Python has a min heap

Mistake 3: Using a heap when sorting is simpler

Mistake 4: Forgetting tuple tie-breaking

Interview Cheat Sheet 🧠

Practice Questions

Easy

Medium

⭐ Must-do interview problems

2. Python `heapq`