Non-Linear Data Structures: Trees and Heaps

TL;DR

You'll understand how trees organize data hierarchically and why that beats linear structures for search and priority tasks. You'll implement a binary search tree and a binary heap from scratch in Python, including insert, delete, and heapify. You'll know exactly when to reach for a BST versus a heap versus a balanced tree in real problems.

1. The Mental Model

A tree is just a chain of decisions: at each node you branch left or right (or into k children) based on some rule. A heap is a tree that only cares about one thing — "who's on top" — not about ordering everything else. Both trade the simplicity of arrays for faster answers to specific questions. A tree is a search structure; a heap is a priority structure — don't ask one to do the other's job.

2. The Core Material

2.1 Binary Search Trees (BST)

Photo by cottonbro studio on Pexels

A BST is a binary tree where, for every node, everything in the left subtree is smaller and everything in the right subtree is larger. This invariant is what makes search fast — you eliminate half the remaining tree at each step, similar to binary search on a sorted array, but without needing contiguous memory or shifting elements on insert.

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

class BST:
    def __init__(self):
        self.root = None

    def insert(self, key):
        self.root = self._insert(self.root, key)

    def _insert(self, node, key):
        if node is None:
            return Node(key)
        if key < node.key:
            node.left = self._insert(node.left, key)
        elif key > node.key:
            node.right = self._insert(node.right, key)
        return node  # duplicates ignored

    def search(self, key):
        node = self.root
        while node:
            if key == node.key:
                return True
            node = node.left if key < node.key else node.right
        return False

    def delete(self, key):
        self.root = self._delete(self.root, key)

    def _delete(self, node, key):
        if node is None:
            return None
        if key < node.key:
            node.left = self._delete(node.left, key)
        elif key > node.key:
            node.right = self._delete(node.right, key)
        else:
            # Found the node to delete
            if node.left is None:
                return node.right
            if node.right is None:
                return node.left
            # Two children: replace with in-order successor
            successor = node.right
            while successor.left:
                successor = successor.left
            node.key = successor.key
            node.right = self._delete(node.right, successor.key)
        return node

The delete case with two children is the part everyone forgets under pressure: you can't just remove the node, you need to pull up a replacement that preserves the ordering invariant. The in-order successor (smallest value in the right subtree) always works because it's bigger than everything on the left and smaller than everything else on the right.

Complexity: search/insert/delete are O(h) where h is the tree's height. For a balanced tree, h = O(log n). For a tree built by inserting sorted data (1, 2, 3, 4, 5...), it degenerates into a linked list and h = O(n). This is the single biggest trap with plain BSTs — they only perform well if the data arrives in a reasonably random order, or if you use a self-balancing variant (AVL, Red-Black) that rotates nodes to cap the height.

2.2 Binary Heaps

Photo by Google DeepMind on Pexels

A heap is a complete binary tree (every level full except possibly the last, filled left to right) satisfying the heap property: in a min-heap, every parent is ≤ its children; in a max-heap, every parent is ≥ its children. Because it's complete, you can store it in a flat array — no pointers needed. For index i (0-indexed): parent is (i-1)//2, left child is 2i+1, right child is 2i+2.

class MinHeap:
    def __init__(self):
        self.heap = []

    def push(self, val):
        self.heap.append(val)
        self._sift_up(len(self.heap) - 1)

    def _sift_up(self, i):
        while i > 0:
            parent = (i - 1) // 2
            if self.heap[i] < self.heap[parent]:
                self.heap[i], self.heap[parent] = self.heap[parent], self.heap[i]
                i = parent
            else:
                break

    def pop(self):
        if not self.heap:
            raise IndexError("pop from empty heap")
        root = self.heap[0]
        last = self.heap.pop()
        if self.heap:
            self.heap[0] = last
            self._sift_down(0)
        return root

    def _sift_down(self, i):
        n = len(self.heap)
        while True:
            left, right = 2 * i + 1, 2 * i + 2
            smallest = i
            if left < n and self.heap[left] < self.heap[smallest]:
                smallest = left
            if right < n and self.heap[right] < self.heap[smallest]:
                smallest = right
            if smallest == i:
                break
            self.heap[i], self.heap[smallest] = self.heap[smallest], self.heap[i]
            i = smallest

push is O(log n) — you bubble the new element up at most h levels. pop is also O(log n) — you move the last element to the root and sift it down. Peeking at the minimum (heap[0]) is O(1), which is the entire point of a heap: fast access to "the best one so far" without sorting everything.

A subtlety worth internalizing: building a heap from n elements all at once is O(n), not O(n log n). If you call push n times, you pay O(n log n) total. But if you use heapify — starting from an arbitrary array and sifting down from the last non-leaf node upward — most nodes are near the bottom (where sift-down is cheap), so the total work sums to O(n). This is why Python's heapq.heapify() exists:

import heapq
data = [9, 3, 7, 1, 8, 2]
heapq.heapify(data)       # O(n), in-place, becomes a valid min-heap array
heapq.heappush(data, 0)   # O(log n)
smallest = heapq.heappop(data)  # O(log n), returns 0

Python only ships a min-heap. For a max-heap, negate values on the way in and out, or push tuples (-priority, item).

2.3 Traversals and Why They Matter

Three depth-first traversal orders answer different questions on a BST:
- In-order (left, node, right) → gives sorted output. Use this to verify a BST is valid or to extract sorted data.
- Pre-order (node, left, right) → gives you the structure/shape, useful for serializing/cloning a tree.
- Post-order (left, right, node) → processes children before the parent, needed for deletion or computing subtree sizes/heights.

def inorder(node, result):
    if node:
        inorder(node.left, result)
        result.append(node.key)
        inorder(node.right, result)

Breadth-first traversal (level order) uses a queue instead of recursion and visits nodes level by level — this is what you want for "shortest path in an unweighted tree" style questions or for printing a tree level by level.

flowchart TD
    A["Need a data structure for..."] --> B{"Fast lookup/search by key<br/>with ordering?"}
    B -->|Yes, balanced insert order| C["BST / AVL / Red-Black Tree<br/>O(log n) search, insert, delete"]
    B -->|No| D{"Need repeated access<br/>to min/max only?"}
    D -->|Yes| E["Binary Heap<br/>O(1) peek, O(log n) push/pop"]
    D -->|No| F{"Need prefix/sorted<br/>string lookups?"}
    F -->|Yes| G["Trie"]
    F -->|No| H["Plain array/list<br/>or hash map"]

3. Worked Example

Problem: Given a stream of numbers [5, 3, 8, 1, 9, 2, 7], find the 3 smallest at any point efficiently, without sorting from scratch each time.

A min-heap is exactly the wrong tool here (it tells you the single minimum fast, but finding "the 3 smallest" means popping 3 times, which destroys the structure for anything beyond that). The right tool for "top-k" problems is actually a max-heap of size k: keep only 3 elements, and when a new number is smaller than the max of those 3, swap it in.

import heapq

def smallest_k(stream, k=3):
    heap = []  # max-heap via negation
    for num in stream:
        if len(heap) < k:
            heapq.heappush(heap, -num)
        elif -num > heap[0]:  # num is smaller than current max of our k
            heapq.heapreplace(heap, -num)
    return sorted(-x for x in heap)

print(smallest_k([5, 3, 8, 1, 9, 2, 7], 3))  # [1, 2, 3]

Walk through it: push 5 → heap [-5]. Push 3 → heap size 2, [-5,-3]. Push 8 → heap size 3, [-8,-3,-5] (max is 8). New number 1: is -1 > -8? Yes → replace, heap becomes size-3 containing {1,3,5}. New number 9: -9 > current max(-top)? 9 is bigger than everything in our set of smallest-so-far, so it's rejected. New number 2: 2 < 5 (current max of the kept set), so it replaces 5. Final kept set: {1,2,3}. Total cost: O(n log k) instead of O(n log n) for a full sort — a real win when k is small and the stream is huge (think: "top 10 trending posts out of a billion events").

This also illustrates the core heap lesson: the heap only guarantees fast access to its extreme element, not fast access to arbitrary ranks. If you needed the 3rd smallest specifically (not "the 3 smallest as a set"), you'd peek heap[0] after this process — that's your answer in O(1).

4. Production Pitfalls & Best Practices

4.1 Real-World Best Practices

• Use heapq from the standard library, don't hand-roll heaps for production code — it's C-optimized and battle-tested.
• For a max-heap, negate on push and pop, or push (-priority, id, item) tuples. Including a unique id avoids TypeError when Python tries to compare two items with equal priority.
• Use heapq.heapify() on existing lists instead of looping heappush — O(n) vs O(n log n).
• Reach for a self-balancing tree (or just use a sorted container / B-tree-backed structure like sortedcontainers.SortedList) when you need guaranteed O(log n) worst case; a plain BST has no such guarantee under adversarial input.
• Store extra metadata in heap entries (e.g., (distance, node) tuples) — this is exactly how Dijkstra's algorithm and A* search are implemented.
• Validate a BST with in-order traversal bounds checking, not by comparing a node only to its immediate children — a common interview mistake that misses violations two levels down.

4.2 Common Bugs & Anti-Patterns

Bug 1 — Comparing only immediate children when validating a BST:

# BAD: misses violations further down the tree
def is_bst(node):
    if not node:
        return True
    if node.left and node.left.key > node.key:
        return False
    if node.right and node.right.key < node.key:
        return False
    return is_bst(node.left) and is_bst(node.right)

# FIXED: pass down valid (low, high) bounds
def is_bst(node, low=float('-inf'), high=float('inf')):
    if not node:
        return True
    if not (low < node.key < high):
        return False
    return (is_bst(node.left, low, node.key) and
            is_bst(node.right, node.key, high))

Bug 2 — Forgetting tuples need a tiebreaker in heapq:

# BAD: crashes if two tasks have equal priority and task objects aren't comparable
heapq.heappush(heap, (priority, task_object))

# FIXED: insert a counter to break ties before Python tries to compare task_object
import itertools
counter = itertools.count()
heapq.heappush(heap, (priority, next(counter), task_object))

Bug 3 — Inserting sorted data into a plain BST and expecting O(log n):

# BAD: inserting 1..1000000 in order turns the BST into a linked list
for i in range(1, 1_000_001):
    bst.insert(i)   # height becomes O(n)!

# FIXED: shuffle insertion order, or use a self-balancing structure
import random
values = list(range(1, 1_000_001))
random.shuffle(values)
for v in values:
    bst.insert(v)   # expected height now O(log n)

Bug 4 — Trying to delete an arbitrary element from a heapq efficiently:

# BAD: heap doesn't support O(log n) arbitrary deletion
heap.remove(target)       # O(n) scan!
heapq.heapify(heap)       # then O(n) again to restore order

# FIXED: lazy deletion — mark as invalid, skip when popped
entry_finder = {}  # maps item -> entry
REMOVED = '<removed>'

def mark_removed(item):
    entry = entry_finder.pop(item)
    entry[-1] = REMOVED  # entry is a list [priority, count, item]

def pop_valid():
    while heap:
        priority, count, item = heapq.heappop(heap)
        if item is not REMOVED:
            del entry_finder[item]
            return item
    raise KeyError('pop from an empty priority queue')

4.3 When to Use What (Decision Matrix)

5. Now Try It

Task: Implement kth_smallest(root, k) that returns the k-th smallest value in a BST using in-order traversal, then modify it to stop early once it hits the k-th element (don't traverse the whole tree if you don't have to).

# Starter — build this tree first, then write kth_smallest
tree = BST()
for v in [8, 3, 10, 1, 6, 14, 4, 7, 13]:
    tree.insert(v)
# kth_smallest(tree.root, 3) should return 4

Write the recursive in-order version first (collect into a list, return result[k-1]). Then rewrite it with an explicit stack-based in-order traversal that returns as soon as you've popped k nodes — this avoids visiting the whole tree when k is small. Test both against k=1 (should give 1), k=5 (should give 7), and k=9 (should give 14, the max).

Success looks like: both versions return [1, 3, 4, 6, 7, 8, 10, 13, 14][k-1] for every k from 1 to 9, and your early-stop version visits noticeably fewer nodes for small k than the full-traversal version.