DSA Interview Prep · Lesson

Python heapq and Max-Heap Tricks

Use heapq.heappush/heappop, negate values to simulate a max-heap, and apply heapq.nlargest/nsmallest for quick top-k queries.

Python heapq and Max-Heap Tricks is a free DSA Interview Prep lesson on CoddyKit — lesson 3 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the DSA Interview Prep learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

Python's heapq Module Overview

Python's heapq module provides a min-heap implemented on top of a regular Python list. Unlike a dedicated heap class, heapq operates on existing lists in-place. The module functions are: heapify to build a heap in O(n), heappush to add an element in O(log n), heappop to remove the minimum in O(log n), and heappushpop / heapreplace for combined efficiency.

import heapq

# heapq operates on plain Python lists
heap = []
heapq.heappush(heap, 5)
heapq.heappush(heap, 2)
heapq.heappush(heap, 8)
heapq.heappush(heap, 1)

print('Heap array:', heap)          # internal array (not sorted!)
print('Peek min:', heap[0])         # O(1) min access
print('Pop min:', heapq.heappop(heap))  # 1
print('Next min:', heap[0])         # 2

# heapify: turn any list into a heap in O(n)
data = [9, 4, 7, 1, 3, 6, 2]
heapq.heapify(data)
print('Heapified:', data, '| min:', data[0])

Max-Heap by Negating Values

Python's heapq only provides a min-heap. To simulate a max-heap, negate all values before pushing and negate again when popping. This works because the heap orders by the stored values, and negating inverts the ordering. Always remember to negate on both sides: negate before push, negate after pop. Forgetting either step is a common interview bug.

import heapq

max_heap = []
for val in [5, 1, 8, 3, 9, 2]:
    heapq.heappush(max_heap, -val)  # negate on push

print('Max-heap internal:', max_heap)  # all negated

# Pop in descending order:
results = []
while max_heap:
    results.append(-heapq.heappop(max_heap))  # negate on pop
print('Sorted descending:', results)  # [9, 8, 5, 3, 2, 1]

# Common pattern: top-k largest
data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
k = 3
heap = []
for x in data:
    heapq.heappush(heap, -x)
print('Top', k, ':', [-heapq.heappop(heap) for _ in range(k)])

heapq.nlargest and nsmallest

heapq.nlargest(k, iterable) and heapq.nsmallest(k, iterable) return the k largest or smallest items. They are O(n log k) — more efficient than full sorting (O(n log n)) when k is much smaller than n. Internally they use a heap of size k. When k is close to n, Python falls back to full sort. Use these for one-shot top-k queries without maintaining a persistent heap.

import heapq

data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 7]

# Top 3 largest:
print(heapq.nlargest(3, data))   # [9, 8, 7]
# Top 3 smallest:
print(heapq.nsmallest(3, data))  # [1, 1, 2]

# With a key function:
words = ['banana', 'apple', 'cherry', 'date', 'elderberry']
print(heapq.nlargest(2, words, key=len))   # ['elderberry', 'banana']
print(heapq.nsmallest(2, words, key=len))  # ['date', 'apple']

# Note: when k ~ n, use sorted() instead:
# sorted(data)[-k:]  or  sorted(data, reverse=True)[:k]

Heap with Tuples for Complex Keys

When heap elements need a custom comparison key, store them as tuples (priority, data). Python's heapq compares tuples element by element, so it first compares priorities. If priorities are equal, it compares the second element — this can cause errors if the data is not comparable. The safest pattern is to include a unique counter as a tiebreaker to avoid ever comparing data elements directly.

import heapq
import itertools

# Pattern: (priority, counter, item)
# Counter ensures unique tiebreaker, avoids comparing items
counter = itertools.count()
heap = []

def push_task(priority, task):
    heapq.heappush(heap, (priority, next(counter), task))

push_task(3, 'low priority task')
push_task(1, 'high priority task')
push_task(2, 'medium priority task')
push_task(1, 'another high priority')

while heap:
    pri, cnt, task = heapq.heappop(heap)
    print(f'P{pri}: {task}')
# Output in priority order: P1, P1, P2, P3

heapq.merge: Merging Sorted Iterables

heapq.merge(*iterables) lazily merges multiple sorted iterables into a single sorted output without loading all data into memory. This is equivalent to a k-way merge using a min-heap of size k and is used in external sort algorithms. It returns an iterator, so elements are produced one at a time — ideal for large datasets or streaming scenarios.

import heapq

# Merge multiple sorted lists efficiently
sorted_lists = [
    [1, 5, 9],
    [2, 6, 8],
    [3, 4, 7]
]

# heapq.merge takes sorted iterables and returns a merged sorted iterator
merged = list(heapq.merge(*sorted_lists))
print('Merged:', merged)  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

# The k-way merge manually (educational version):
def merge_k_sorted(lists):
    heap = []
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))
    result = []
    while heap:
        val, list_idx, elem_idx = heapq.heappop(heap)
        result.append(val)
        if elem_idx + 1 < len(lists[list_idx]):
            heapq.heappush(heap, (lists[list_idx][elem_idx+1], list_idx, elem_idx+1))
    return result

print('Manual k-way:', merge_k_sorted(sorted_lists))

Lazy Deletion Pattern for Heaps

When you need to remove arbitrary elements from a heap but don't know their index, use lazy deletion: mark elements as deleted in a separate set, then skip them when popping. This is O(log n) amortised and avoids the complexity of tracking indices. It is the standard approach in Dijkstra's algorithm with duplicate entries and task-scheduler simulations.

import heapq

class LazyHeap:
    def __init__(self):
        self._heap = []
        self._removed = set()

    def push(self, task):
        heapq.heappush(self._heap, task)

    def remove(self, task):
        self._removed.add(task)  # mark as removed

    def pop(self):
        while self._heap:
            task = heapq.heappop(self._heap)
            if task not in self._removed:
                return task
        return None

lh = LazyHeap()
for t in [5, 1, 8, 3, 2]:
    lh.push(t)
lh.remove(1)  # 'delete' 1 lazily
lh.remove(8)  # 'delete' 8 lazily
results = [lh.pop() for _ in range(3)]
print(results)  # [2, 3, 5] -- 1 and 8 skipped

Kth Largest Element in a Stream

Kth Largest Element in a Stream (LeetCode #703) maintains a min-heap of size k. The root of the heap is always the kth largest element seen so far. When a new number arrives: push it, and if the heap exceeds size k, pop the minimum. The root is always the kth largest because there are exactly k-1 elements larger than it in the heap.

import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)  # remove smallest
        return self.heap[0]  # kth largest = root of min-heap

# k=3, initial=[4,5,8,2]
kl = KthLargest(3, [4, 5, 8, 2])
print(kl.add(3))   # 4 (top 3: 8,5,4 -- kth=4)
print(kl.add(5))   # 5 (top 3: 8,5,5 -- kth=5)
print(kl.add(10))  # 5 (top 3: 10,8,5 -- kth=5)
print(kl.add(9))   # 8 (top 3: 10,9,8 -- kth=8)

Find K Pairs with Smallest Sum

Find K pairs with smallest sums (LeetCode #373) uses a min-heap to generate pairs in order. Start with all pairs (nums1[0], nums2[j]) for each j. Pop the minimum, and for the popped pair (nums1[i], nums2[j]), push (nums1[i+1], nums2[j]) — the next candidate from the same nums2 column. This is a common pattern for ordered pair/product generation with a heap.

import heapq

def k_smallest_pairs(nums1, nums2, k):
    if not nums1 or not nums2:
        return []
    heap = []
    # Initialize with pairs (nums1[0], nums2[j])
    for j in range(min(k, len(nums2))):
        heapq.heappush(heap, (nums1[0] + nums2[j], 0, j))
    result = []
    while heap and len(result) < k:
        total, i, j = heapq.heappop(heap)
        result.append([nums1[i], nums2[j]])
        if i + 1 < len(nums1):
            heapq.heappush(heap, (nums1[i+1] + nums2[j], i+1, j))
    return result

print(k_smallest_pairs([1,7,11], [2,4,6], 3))
# [[1,2], [1,4], [1,6]]

Task Scheduler with a Max-Heap

Task Scheduler (LeetCode #621) asks for the minimum time to schedule n tasks with a cooling period of n intervals between same tasks. Use a max-heap of task frequencies: at each time step, pick the most frequent available task, decrement its count, and put it on cooldown. Process k=n+1 tasks per cycle (or fill with idle time). This greedy approach with a max-heap gives the optimal answer.

import heapq
from collections import Counter

def least_interval(tasks, n):
    freq = Counter(tasks)
    heap = [-f for f in freq.values()]  # max-heap (negated)
    heapq.heapify(heap)
    time = 0
    while heap:
        cycle = n + 1
        temp = []
        for _ in range(cycle):
            if heap:
                temp.append(heapq.heappop(heap))
        for f in temp:
            if f + 1 < 0:  # still tasks remaining
                heapq.heappush(heap, f + 1)
        # Add full cycle or remaining tasks if queue empty
        time += cycle if heap else len(temp)
    return time

print(least_interval(['A','A','A','B','B','B'], 2))  # 8
print(least_interval(['A','A','A','B','B','B'], 0))  # 6

Heap in Dijkstra's Algorithm

The priority queue in Dijkstra's algorithm is implemented with a min-heap. Store tuples (distance, node) and always process the nearest unvisited node first. When you pop a node with a distance greater than its currently known shortest path (a stale entry from lazy deletion), skip it. This avoids the need for a decrease-key operation and keeps the implementation simple while maintaining O((V + E) log V) complexity.

import heapq

def dijkstra(graph, start):
    dist = {node: float('inf') for node in graph}
    dist[start] = 0
    heap = [(0, start)]  # (distance, node)
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist[u]:   # stale entry, skip
            continue
        for v, w in graph[u]:
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                heapq.heappush(heap, (dist[v], v))
    return dist

graph = {
    'A': [('B', 4), ('C', 1)],
    'B': [('D', 1)],
    'C': [('B', 2), ('D', 5)],
    'D': []
}
print(dijkstra(graph, 'A'))  # {'A':0,'B':3,'C':1,'D':4}

Reorganize String with a Max-Heap

Reorganize String (LeetCode #767) asks you to rearrange a string so no two adjacent characters are the same. Use a max-heap of (-frequency, char). At each step, pop the most frequent character. If the previous character is the same as the most frequent, pop the second-most frequent instead. This greedy approach ensures the most constrained character is placed as early as possible.

import heapq
from collections import Counter

def reorganize_string(s):
    freq = Counter(s)
    heap = [(-f, c) for c, f in freq.items()]
    heapq.heapify(heap)
    result = []
    prev_freq, prev_char = 0, ''
    while heap:
        freq, char = heapq.heappop(heap)
        result.append(char)
        # Push back the previous character if still remaining
        if prev_freq < 0:
            heapq.heappush(heap, (prev_freq, prev_char))
        prev_freq, prev_char = freq + 1, char  # decrement freq (less negative)
    result_str = ''.join(result)
    # Verify no adjacent duplicates
    return result_str if len(result_str) == len(s) else ''

print(reorganize_string('aab'))   # 'aba'
print(reorganize_string('aaab'))  # '' (impossible)

Quick Check

Test your understanding of Data Structures & Algorithms — Coding Interview Prep concepts from this lesson.

Lesson Recap

In this lesson you learned: Python's heapq module API including heapify, heappush, heappop, nlargest, nsmallest, and merge, max-heap simulation by negating values, and common heap interview patterns including top-k streaming, kth largest in a stream, task scheduler, and Dijkstra. Next up we tackle median from data stream and k-way merge.

Frequently asked questions

Is the “Python heapq and Max-Heap Tricks” lesson free?

Yes — the full text of “Python heapq and Max-Heap Tricks” is free to read here on the web, and the DSA Interview Prep course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the DSA Interview Prep course, upgrade to CoddyKit PRO.

What will I learn in “Python heapq and Max-Heap Tricks”?

Use heapq.heappush/heappop, negate values to simulate a max-heap, and apply heapq.nlargest/nsmallest for quick top-k queries. You practise DSA Interview Prep with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start DSA Interview Prep?

No prior experience is required. DSA Interview Prep on CoddyKit is structured for beginners through advanced learners; this is — lesson 3 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Python heapq and Max-Heap Tricks” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this DSA Interview Prep lesson?

Yes. Every DSA Interview Prep lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

← Back to DSA Interview Prep