heapq and bisect

Part of the free LogicWiz Python course, Module 12: Iteration and the standard toolkit. The examples and practice run in the app.

At lunch rush, Shankar bhaiya serves teachers first, then students in token order. And the result portal needs the top 3 scores out of 2 lakh. Sorting everything each time is slow; a heap is built for this.

You'll learn to

  • Use `heapq` as a min-heap: push, pop and peek the smallest
  • Get the top k with `nlargest`/`nsmallest`, and build a max-heap by negating
  • Build a priority queue with tuples
  • Search and insert in a sorted list with `bisect`

Learn

A heap keeps the smallest on top

A heap is a list arranged so that h[0] is always the smallest item. Python's heapq module works on a plain list:

  • heappush(h, x): add, O(log n)
  • heappop(h): remove and return the smallest, O(log n)
  • h[0]: peek at the smallest, O(1)
  • heapify(list): turn a list into a heap in place, O(n)

The rest of the list is not sorted. Only the top is guaranteed.

Push and pop

import heapq

h = []
for price in [30, 10, 20]:
    heapq.heappush(h, price)
print(h[0])
print(heapq.heappop(h))
print(heapq.heappop(h))
print(h)

Top k with nlargest and nsmallest

import heapq

scores = [72, 95, 64, 88, 91, 59]
print(heapq.nlargest(3, scores))
print(heapq.nsmallest(2, scores))

menu = [("chai", 10), ("dosa", 40), ("samosa", 15), ("thali", 90)]
print(heapq.nlargest(2, menu, key=lambda m: m[1]))

For a small k out of a huge list, this is O(n log k), faster than sorting everything (O(n log n)).

A max-heap: store negatives

import heapq

bills = [250, 80, 600, 120]
h = [-b for b in bills]       # biggest becomes most negative
heapq.heapify(h)
print(-heapq.heappop(h))      # flip back on the way out
print(-heapq.heappop(h))

A priority queue with tuples

import heapq

queue = []
order_no = 0
for name, priority in [("Riya", 2), ("Mr Rao", 1), ("Kabir", 2), ("Ms Iyer", 1)]:
    order_no += 1
    heapq.heappush(queue, (priority, order_no, name))

while queue:
    priority, _, name = heapq.heappop(queue)
    print(priority, name)

Tuples compare item by item: priority first, then the arrival number breaks ties so equal priorities are served in order.

bisect: binary search on a sorted list

bisect finds where a value belongs in a sorted list in O(log n):

  • bisect_left(a, x): the first position where x could go (before any equal items)
  • bisect_right(a, x) (or bisect): the position after any equal items
  • insort(a, x): insert x and keep the list sorted

The list must already be sorted, or the answers are garbage.

Grades from cut-offs, and counting

import bisect

cutoffs = [40, 60, 75, 90]
grades = "FDCBA"
def grade(score):
    return grades[bisect.bisect_right(cutoffs, score)]
print([grade(s) for s in [35, 40, 74, 75, 99]])

marks = [40, 55, 55, 55, 70]
print(bisect.bisect_left(marks, 55), bisect.bisect_right(marks, 55))

bisect.insort(marks, 60)
print(marks)

Two heap surprises:

  • print(h) is not sorted. Only h[0] is guaranteed to be the smallest. To get everything in order, pop repeatedly (or just use sorted).
  • heapq is only a min-heap. For the largest first, push -x and negate when you pop. For tuples, negate the number field: (-amount, name).

"Find the k largest elements", "kth largest", "merge k sorted lists" and "meeting rooms / task scheduling" are standard heap interview questions. Say the complexity: building a heap O(n), each push/pop O(log n), top k with a size-k heap O(n log k).

In one line

heapq is a min-heap on a list (h[0] smallest, push/pop O(log n)); negate for a max-heap, use nlargest(k, …) for top k; bisect searches and insort inserts into a sorted list in O(log n).

Interview corner

How do you find the k largest elements in a big list efficiently?

heapq.nlargest(k, nums), or keep a min-heap of size k: push each number and pop when the size exceeds k; the heap then holds the k largest. That is O(n log k) time and O(k) memory, better than sorting (O(n log n)) when k is small.

Python only has a min-heap. How do you make a max-heap?

Push negated values (heappush(h, -x)) and negate again when popping (-heappop(h)). For tuples, negate the numeric key: (-score, name). For objects, wrap them in a class whose __lt__ is reversed.

Why add a counter to heap entries like (priority, count, item)?

When priorities tie, tuples compare the next field. A rising counter keeps equal priorities in insertion order (stable) and stops Python from comparing the items themselves, which may not support < at all (e.g. dicts) and would raise a TypeError.

What is bisect used for?

Binary search on a sorted list: bisect_left/bisect_right find insertion points in O(log n), useful for grade or tax slabs, counting values in a range, and keeping a list sorted with insort (the search is O(log n) but the insert shifts items, O(n)).