Trees and Heaps in Python: Traversals and heapq in Practice

A disk-usage report has to answer "how big is this folder, counting everything inside it?" Folders contain folders, which contain folders, to any depth. A flat list cannot describe that, but a tree can, and the natural way to process a tree is to treat each child as a smaller version of the same problem. This lesson starts there, then moves to the heap, a tree stored in an array that always keeps its smallest item at the top.

A tree is a node plus child trees

Each Folder node stores its own bytes and a list of children. Because a subtree has the same shape as the whole tree, a recursive function writes itself: a node's total is its own size plus the totals of its children.

from collections import deque


class Folder:
    def __init__(self, name, size=0, children=None):
        self.name = name
        self.size = size                # bytes stored directly in this folder
        self.children = children or []


root = Folder("project", 10, [
    Folder("src", 40, [Folder("api", 120), Folder("ui", 200, [Folder("icons", 75)])]),
    Folder("docs", 30),
    Folder("tests", 90, [Folder("fixtures", 300)]),
])


def total_size(node):
    return node.size + sum(total_size(child) for child in node.children)


def depth_first(node, depth=0):
    yield "  " * depth + node.name
    for child in node.children:
        yield from depth_first(child, depth + 1)


def breadth_first(node):
    queue = deque([node])
    while queue:
        current = queue.popleft()
        yield current.name
        queue.extend(current.children)


print("total:", total_size(root))
print("\n".join(depth_first(root)))
print("BFS:", list(breadth_first(root)))

total: 865
project
  src
    api
    ui
      icons
  docs
  tests
    fixtures
BFS: ['project', 'src', 'docs', 'tests', 'api', 'ui', 'fixtures', 'icons']

The two traversals visit the same eight folders in different orders. Depth-first goes all the way down one branch before backing up (recursion supplies the stack from the previous lesson). Breadth-first visits everything one level at a time using a queue, which is why it finds the shallowest match first.

Binary search trees and why insertion order matters

A binary search tree keeps smaller keys in the left subtree and larger ones in the right, so a lookup discards half the remaining tree at each step. Reading it in-order returns keys sorted. The catch is that the shape follows the insertion order.

class Node:
    def __init__(self, key):
        self.key, self.left, self.right = key, None, None


def insert(root, key):
    if root is None:
        return Node(key)
    if key < root.key:
        root.left = insert(root.left, key)
    elif key > root.key:
        root.right = insert(root.right, key)
    return root


def in_order(node):
    if node:
        yield from in_order(node.left)
        yield node.key
        yield from in_order(node.right)


def height(node):
    return 0 if node is None else 1 + max(height(node.left), height(node.right))


root = None
for key in [50, 30, 70, 20, 40, 60, 80]:
    root = insert(root, key)
print(list(in_order(root)), "height", height(root))

lopsided = None
for key in range(1, 8):          # keys arrive already sorted
    lopsided = insert(lopsided, key)
print(list(in_order(lopsided)), "height", height(lopsided))

[20, 30, 40, 50, 60, 70, 80] height 3
[1, 2, 3, 4, 5, 6, 7] height 7

Same seven keys, same sorted output, but heights 3 and 7. The second tree is a linked list in disguise, and a search in it is O(n). Production search trees (red-black, AVL, B-trees) rebalance themselves to keep the height near log n. In everyday Python you rarely hand-write one; a dict covers key lookups and bisect on a sorted list covers ordered lookups. Databases use B-trees for indexes, which is covered in Database Indexing.

Heaps: always know the smallest item

A min-heap is a tree with one rule: every parent is less than or equal to its children. That does not sort the whole thing, but the smallest item is always at the root, and heapq stores the tree in a plain list. For the item at index i, the parent is at (i - 1) // 2 and the children are at 2*i + 1 and 2*i + 2.

list:   [1, 3, 2, 9, 4]
index:   0  1  2  3  4

        1            parent of idx 3 and 4 is idx 1
       / \           parent of idx 1 and 2 is idx 0
      3   2
     / \
    9   4

Pushing and popping are O(log n). A print shop that always wants the most urgent job next is the classic fit. Ties must be broken by something comparable, otherwise Python would try to compare the job payloads themselves, so an arrival counter goes in the middle of the tuple.

import heapq
from itertools import count


class JobBoard:
    """Print-shop jobs: lowest priority number first, ties served in arrival order."""

    def __init__(self):
        self._heap = []
        self._arrival = count()          # tie-breaker so we never compare job names

    def submit(self, priority, job):
        heapq.heappush(self._heap, (priority, next(self._arrival), job))

    def next_job(self):
        priority, _, job = heapq.heappop(self._heap)
        return priority, job

    def __len__(self):
        return len(self._heap)


board = JobBoard()
board.submit(2, "banner")
board.submit(1, "wedding invites")
board.submit(2, "flyers")
board.submit(3, "stickers")
board.submit(1, "rush business cards")

while board:
    print(board.next_job())

(1, 'wedding invites')
(1, 'rush business cards')
(2, 'banner')
(2, 'flyers')
(3, 'stickers')

Top-k without sorting everything

To find the three largest of millions of file sizes, sorting is more work than needed. Keep a min-heap of size k: its root is the weakest of the current best k, so a new value only matters if it beats the root. Memory stays at k items.

import heapq


def top_k_largest(sizes, k):
    """Keep only k items in memory: a min-heap whose root is the weakest of the current top k."""
    heap = []
    for name, size in sizes:
        if len(heap) < k:
            heapq.heappush(heap, (size, name))
        elif size > heap[0][0]:
            heapq.heapreplace(heap, (size, name))     # pop smallest, push new, one step
    return sorted(heap, reverse=True)


files = [("a.log", 420), ("b.mp4", 9100), ("c.txt", 3), ("d.zip", 5300),
         ("e.png", 880), ("f.iso", 7700), ("g.csv", 640)]

print(top_k_largest(files, 3))
print(heapq.nlargest(3, files, key=lambda f: f[1]))      # same answer, built in

data = [9, 4, 7, 1, 8, 2]
heapq.heapify(data)
print(data)          # a valid heap, not a sorted list
print(data[0])       # smallest is always at index 0

[(9100, 'b.mp4'), (7700, 'f.iso'), (5300, 'd.zip')]
[('b.mp4', 9100), ('f.iso', 7700), ('d.zip', 5300)]
[1, 4, 2, 9, 8, 7]
1

Why is heapq.heapreplace one step instead of two? A push adds the item at the end and lets it float up past larger parents; a pop moves the last item to the root and lets it sink below smaller children. Each move touches one level of the tree, and a tree of n items has about log n levels, which is where O(log n) comes from. The top-k loop therefore costs O(n log k): n items, each at worst one trip through a heap that never holds more than k.

Common mistakes

• Treating the heap list as sorted. Symptom: data[1] is not the second smallest (above it is 4, but the second smallest of the data is 2). Only data[0] is guaranteed; pop repeatedly or use sorted for the full order.
• Pushing tuples whose tie-break element cannot be compared. Symptom: TypeError: '<' not supported between instances of 'dict' and 'dict' when two priorities are equal. Insert a counter as the second tuple element.
• Expecting `heapq` to be a max-heap. It is a min-heap only. Negate numbers on the way in and out, or push (-priority, ...).
• Recursing on a very deep tree. Symptom: RecursionError: maximum recursion depth exceeded. A chain-shaped tree (like the lopsided BST at scale) needs an explicit stack or queue.
• Building a BST from sorted input and expecting log n. Symptom: lookups slow down linearly. Use a balanced structure or shuffle the keys.

How this shows up in real projects and interviews

Trees are everywhere you stop looking for them: file systems, HTML and JSON documents, org charts, comment threads, database indexes. Heaps drive schedulers, "k most frequent" features, event simulations and merging sorted streams. Interviews ask for traversal orders, "height of a tree", "level order output" and "k largest elements", and a strong answer names the structure first (queue for levels, bounded min-heap for top-k) and then states the cost, O(n log k) for the heap version. If the follow-up is "what if the data is already sorted?", revisit Sorting Algorithms, and for array-shaped patterns see Arrays and Strings.