Heaps, Tries and Disjoint Sets

The binary heap

A priority queue has to answer one question again and again: which waiting item is the smallest? A sorted list answers it instantly but is slow to insert into, because everything after the new item has to shift along. An unsorted list inserts instantly but needs a full scan to find the smallest. A binary heap is the compromise: it keeps the items only partly ordered, which is just enough.

A heap is a complete binary tree: every level is full except possibly the last, which fills from the left. In a min-heap it also obeys the heap property:

Because the tree is complete, it needs no pointers. It is stored in a plain list, level by level, and the position of a node tells you where its relatives are:

heap = [1, 3, 8, 5, 4, 9]
i = 1
print("node", heap[i])
print("left child", heap[2 * i + 1])
print("right child", heap[2 * i + 2])
print("parent", heap[(i - 1) // 2])

Adding an item: sift up

Put the new item at the end, so the tree stays complete. It may now be smaller than its parent, which breaks the property, so swap it with its parent, and repeat until it is in place or reaches the root. This is sifting up.

Removing the smallest: sift down

The smallest is the root. Take it out, move the last item into the root's place so the tree stays complete, then swap that item downwards with its smaller child until the property holds again. This is sifting down.

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

    def push(self, value):
        self.items.append(value)
        child = len(self.items) - 1
        while child > 0:
            parent = (child - 1) // 2
            if self.items[parent] <= self.items[child]:
                break
            self.items[parent], self.items[child] = self.items[child], self.items[parent]
            child = parent

    def pop(self):
        top = self.items[0]
        last = self.items.pop()
        if self.items:
            self.items[0] = last
            parent = 0
            size = len(self.items)
            while True:
                left, right = 2 * parent + 1, 2 * parent + 2
                smallest = parent
                if left < size and self.items[left] < self.items[smallest]:
                    smallest = left
                if right < size and self.items[right] < self.items[smallest]:
                    smallest = right
                if smallest == parent:
                    break
                self.items[parent], self.items[smallest] = self.items[smallest], self.items[parent]
                parent = smallest
        return top

h = MinHeap()
for value in [5, 3, 8, 1, 9, 2]:
    h.push(value)
    print("after pushing", value, "->", h.items)

print("removed", h.pop(), "->", h.items)
print("removed", h.pop(), "->", h.items)

Notice that the list is not sorted, and does not need to be. Only the smallest is guaranteed to be at the front.

The cost

Both operations move an item at most the height of the tree, and a complete tree of n items has a height of about log₂ n. Looking at the smallest is just items[0].

Heap sort and the top k problem

Heap sort

Push everything into a heap, then pop everything out. The items leave smallest first, so they come out sorted. Each of the n pushes and n pops costs O(log n), so the whole sort is O(n log n), the same as merge sort, and unlike quick sort it has no bad case.

import heapq

def heap_sort(values):
    heap = list(values)
    heapq.heapify(heap)           # turns a list into a heap in O(n)
    return [heapq.heappop(heap) for _ in range(len(heap))]

print(heap_sort([9, 4, 7, 1, 8, 2, 6]))

heapq is Python's built in heap. It works on an ordinary list and keeps a min-heap. heapify builds a heap from a list in O(n), faster than pushing items one at a time.

The top k problem

How would you find the five highest scores out of a hundred thousand? Sorting everything costs O(n log n) and holds every score in memory. A heap does better. Keep a small min-heap of the best k seen so far. Its root is the weakest of the best, so a new score only matters if it beats the root, and then it replaces it.

import heapq

def top_k(stream, k):
    best = []
    for score in stream:
        if len(best) < k:
            heapq.heappush(best, score)
        elif score > best[0]:
            heapq.heapreplace(best, score)
    return sorted(best, reverse=True)

scores = [(i * 7919) % 100003 for i in range(100000)]
print(top_k(scores, 5))
print(sorted(scores, reverse=True)[:5] == top_k(scores, 5))

The heap never holds more than five items, however long the stream. That makes it usable on data too big to fit in memory, such as a log file read one line at a time. Python wraps the idea up as heapq.nlargest(k, data) and heapq.nsmallest(k, data).

Priority with a payload

A heap compares whatever it stores. Store a tuple and it compares the first item, then the second if there is a tie, which is how the task queue in the stacks and queues topic works.

import heapq

jobs = []
heapq.heappush(jobs, (2, "send newsletter"))
heapq.heappush(jobs, (1, "restart server"))
heapq.heappush(jobs, (2, "back up database"))

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

Tries: searching by prefix

A hash table finds a whole key fast. It cannot answer "which words start with car?" without checking every key. A trie (pronounced "try", from re*trie*val) is a tree built for exactly that question.

Each node stands for one letter, and the path from the root spells a prefix. Words that share a beginning share the nodes for it, so car, card, care and cart all hang from the same c, a, r. A marker records where a real word ends.

class Trie:
    def __init__(self):
        self.root = {}

    def insert(self, word):
        node = self.root
        for letter in word:
            node = node.setdefault(letter, {})
        node["$"] = True                 # a word ends here

    def find(self, prefix):
        node = self.root
        for letter in prefix:
            if letter not in node:
                return None
            node = node[letter]
        return node

    def contains(self, word):
        node = self.find(word)
        return node is not None and "$" in node

    def suggest(self, prefix):
        found = []

        def walk(node, text):
            for letter in sorted(node):
                if letter == "$":
                    found.append(text)
                else:
                    walk(node[letter], text + letter)

        start = self.find(prefix)
        if start is not None:
            walk(start, prefix)
        return found

words = Trie()
for w in ["car", "card", "care", "careful", "cart", "cat", "do", "dog"]:
    words.insert(w)

print(words.contains("care"), words.contains("ca"))
print(words.suggest("car"))
print(words.suggest("d"))
print(words.suggest("x"))

contains("ca") is False even though the path exists, because no word ends there. That is the job of the "$" marker.

Why it is quick

Looking up a word takes one step per letter, so the cost depends on the length of the word, not on how many words are stored. A trie of ten words and a trie of ten million both find care in four steps. Autocomplete uses this: follow the letters typed so far, then collect everything below that node.

Disjoint sets and union-find

A social network wants to know whether two people are connected through a chain of friends. A game wants to know whether two rooms are joined. A network engineer wants to know whether adding a cable would join two groups that are already connected. All are asking: are these two items in the same group?, with groups that keep merging.

A disjoint set structure, usually called union-find, answers it with two operations:

• find(x) returns the representative (the root) of the group that holds x
• union(a, b) merges the groups that hold a and b

Each group is a tree stored in a list. parent[x] points at x's parent, and a root points at itself. Two items are in the same group if and only if find gives them the same root.

class DisjointSets:
    def __init__(self, n):
        self.parent = list(range(n))     # at the start everyone is alone
        self.size = [1] * n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])    # path compression
        return self.parent[x]

    def union(self, a, b):
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return False                 # already in the same group
        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a     # hang the smaller tree under the bigger
        self.size[root_a] += self.size[root_b]
        return True

friends = DisjointSets(8)
for a, b in [(0, 1), (2, 3), (1, 2), (5, 6)]:
    friends.union(a, b)

print(friends.find(0) == friends.find(3))     # 0 - 1 - 2 - 3
print(friends.find(0) == friends.find(5))
groups = {friends.find(x) for x in range(8)}
print(len(groups), "groups")

Two small ideas make it fast

• Union by size. Hang the smaller tree under the root of the bigger one, so trees stay shallow.
• Path compression. While find walks up to the root, it points every item it passes straight at the root. The next find from any of them takes one step.

parent = [1, 2, 3, 4, 5, 6, 7, 8, 9, 9]      # a chain: 0 -> 1 -> 2 -> ... -> 9

def find(x):
    steps = 0
    path = []
    while parent[x] != x:
        path.append(x)
        x = parent[x]
        steps += 1
    for p in path:
        parent[p] = x                 # compress the path
    return x, steps

print(find(0))       # the first search walks the whole chain
print(find(0))       # now it is one step
print(find(4))

Together these make find and union run in close to constant time, however large the sets become.

Minimum spanning trees

A spanning tree of a graph joins every vertex using just enough edges to stay connected and no cycles: always one fewer edge than vertices. A minimum spanning tree (MST) is the spanning tree whose edges have the smallest possible total weight.

It is the answer to: what is the cheapest way to connect all of these places? Laying cable between offices, pipes between houses, or roads between villages. It is not the same as Dijkstra's algorithm, which finds the cheapest route from one vertex to another.

Both algorithms below are greedy: they repeatedly take the cheapest choice available, and for this problem that is provably enough.

Kruskal's algorithm

Sort all edges by weight. Take them in order, and accept an edge unless it would form a cycle. Union-find answers "would it form a cycle?" in one call: if both ends are already in the same group, the edge is a cycle.

class DisjointSets:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True

vertices = ["A", "B", "C", "D", "E", "F", "G"]
edges = [("A", "B", 7), ("A", "D", 5), ("B", "C", 8), ("B", "D", 9), ("B", "E", 7),
         ("C", "E", 5), ("D", "E", 15), ("D", "F", 6), ("E", "F", 8), ("E", "G", 9),
         ("F", "G", 11)]

def kruskal(vertices, edges):
    number = {v: i for i, v in enumerate(vertices)}
    groups = DisjointSets(len(vertices))
    chosen, total = [], 0
    for a, b, weight in sorted(edges, key=lambda e: e[2]):
        if groups.union(number[a], number[b]):
            print("accept", a + b, weight)
            chosen.append((a, b, weight))
            total += weight
        else:
            print("reject", a + b, weight, "(would make a cycle)")
    return chosen, total

tree, cost = kruskal(vertices, edges)
print(len(tree), "edges, total weight", cost)

Prim's algorithm

Prim grows a single tree outwards. Start at any vertex. Repeatedly take the cheapest edge that leads out of the tree to a new vertex. A priority queue holds the candidate edges, so this is the same shape as Dijkstra's algorithm.

import heapq

edges = [("A", "B", 7), ("A", "D", 5), ("B", "C", 8), ("B", "D", 9), ("B", "E", 7),
         ("C", "E", 5), ("D", "E", 15), ("D", "F", 6), ("E", "F", 8), ("E", "G", 9),
         ("F", "G", 11)]

def prim(edges, start):
    graph = {}
    for a, b, weight in edges:
        graph.setdefault(a, []).append((weight, a, b))
        graph.setdefault(b, []).append((weight, b, a))
    inside = {start}
    queue = list(graph[start])
    heapq.heapify(queue)
    total = 0
    while queue and len(inside) < len(graph):
        weight, a, b = heapq.heappop(queue)
        if b in inside:
            continue
        inside.add(b)
        total += weight
        print("add", b, "via", a + b, weight)
        for edge in graph[b]:
            if edge[2] not in inside:
                heapq.heappush(queue, edge)
    return total

print("total weight", prim(edges, "A"))

Different algorithms, same total: when edge weights tie, there can be more than one minimum spanning tree, but they all cost the same.

Key terms you must be able to define

Heap

A complete binary tree stored in a list, in which every parent is ordered with respect to its children.

Heap property

In a min-heap, every parent is less than or equal to its children, so the smallest item is at the root.

Sift up

Swapping a newly added item with its parent until the heap property holds.

Sift down

Swapping the item moved to the root with its smaller child until the heap property holds.

Heap sort

Sorting by building a heap and repeatedly removing the smallest item. O(n log n).

Trie

A tree in which each node is a letter and each path from the root spells a prefix, used to search by prefix.

Disjoint set (union-find)

A structure that keeps items in non overlapping groups, supporting find (which group?) and union (merge two groups).

Path compression

Making every item passed during a find point directly at the root, so later finds are quicker.

Spanning tree

A tree that connects every vertex of a graph using the fewest edges, with no cycles.

Minimum spanning tree

The spanning tree with the smallest total edge weight.

Greedy algorithm

An algorithm that makes the best looking choice at each step and never reconsiders it.

Flashcards

Say the definition out loud before you flip the card. Retrieving it from memory is what builds recall, reading it again is not.

Reach the top grade on this topic

Heaps and trees of this kind are an extension beyond the specification list, so they appear as unseen structures. The marks go to reasoning, not memory.

• State the property in one sentence: every parent is at most its children, so the minimum is at the root
• When tracing an insertion, write the array after each swap and name the parent index you compared with
• Compare costs with the alternatives: a heap gives O(log n) for both add and remove, where each list gives O(n) for one of them
• For a trie, say the cost depends on word length, not on how many words are stored
• For a spanning tree, say how many edges the answer has: one fewer than the number of vertices

Mistakes that cost marks here

index:  0  1  2  3  4  5
value:  2  5  4  9  7  8

def find(parent, x):
    if parent[x] != x:
        parent[x] = find(parent, parent[x])
    return parent[x]

parent = [1, 2, 3, 3]
print(find(parent, 0))
print(parent)