Heaps & priority queues

Always know the smallest (or largest) item in O(1), add and remove in O(log n): how heaps work inside, container/heap, and top-k problems.

The heap property#

A (min-)heap is a binary tree where every parent is smaller than or equal to its children. So the smallest element is always at the root. Siblings are not ordered relative to each other; a heap is much looser than a sorted array, which is why it is cheaper to maintain.

It is also complete: every level is full except possibly the last, filled left to right. That lets us store it in a plain slice with no pointers:

        1              index:  0  1  2  3  4  5
      /   \            value: [1, 3, 2, 7, 4, 5]
     3     2
    / \   /            children of i: 2i+1, 2i+2
   7   4 5             parent of i:   (i-1)/2

Push: add at the end, sift up#

Append the new value, then while it is smaller than its parent, swap them.

Push 0 into [1, 3, 2, 7, 4, 5]:

step	slice	new value at index	parent	swap?
append	[1 3 2 7 4 5 0]	6	index 2 (value 2)	yes
	[1 3 0 7 4 5 2]	2	index 0 (value 1)	yes
	[0 3 1 7 4 5 2]	0	none	done

At most one swap per level: O(log n).

Run the trace

A heap written by hand so every swap is visible. The index math from the diagram is all it needs:

siftup_trace.go

package main
 
import "fmt"
 
func push(h []int, v int) []int {
	h = append(h, v)
	i := len(h) - 1
	for i > 0 {
		p := (i - 1) / 2
		if h[p] <= h[i] {
			fmt.Printf("  value %d at index %d >= parent value %d: done\n", h[i], i, h[p])
			break
		}
		h[p], h[i] = h[i], h[p]
		fmt.Printf("  swap with parent at %d -> %v\n", p, h)
		i = p
	}
	return h
}
 
func main() {
	h := []int{1, 3, 2, 7, 4, 5}
	fmt.Println("push 0 into", h)
	h = push(h, 0)
	fmt.Println("push 6 into", h)
	h = push(h, 6)
	fmt.Println("heap:", h)
}

Output

push 0 into [1 3 2 7 4 5]
  swap with parent at 2 -> [1 3 0 7 4 5 2]
  swap with parent at 0 -> [0 3 1 7 4 5 2]
push 6 into [0 3 1 7 4 5 2]
  swap with parent at 3 -> [0 3 1 6 4 5 2 7]
  value 6 at index 3 >= parent value 3: done
heap: [0 3 1 6 4 5 2 7]

Pushing 6 needs only one swap: it is smaller than its first parent (7) but not smaller than the next one (3). Most pushes stop after a swap or two, long before reaching the root.

Pop: take the root, sift down#

Remove the root (the minimum), move the last element to the root, then while it is larger than its smaller child, swap with that child.

Pop from [0, 3, 1, 7, 4, 5, 2]:

step	slice	moving value at	smaller child	swap?
take 0, move last (2) to root	[2 3 1 7 4 5]	0	1 (index 2)	yes
	[1 3 2 7 4 5]	2	5 (index 5)	no: 2 < 5, done

Also O(log n).

Run the trace

siftdown_trace.go

package main
 
import "fmt"
 
func pop(h []int) (int, []int) {
	top := h[0]
	last := len(h) - 1
	h[0] = h[last] // move the last element to the root
	h = h[:last]
	fmt.Printf("  took %d, moved %d to root -> %v\n", top, h[0], h)
	i := 0
	for {
		small := i
		for _, c := range []int{2*i + 1, 2*i + 2} {
			if c < len(h) && h[c] < h[small] {
				small = c
			}
		}
		if small == i {
			fmt.Println("  no smaller child: done")
			break
		}
		h[i], h[small] = h[small], h[i]
		fmt.Printf("  swap with child at %d -> %v\n", small, h)
		i = small
	}
	return top, h
}
 
func main() {
	h := []int{0, 3, 1, 7, 4, 5, 2}
	for len(h) > 4 {
		fmt.Println("pop from", h)
		_, h = pop(h)
	}
}

Output

pop from [0 3 1 7 4 5 2]
  took 0, moved 2 to root -> [2 3 1 7 4 5]
  swap with child at 2 -> [1 3 2 7 4 5]
  no smaller child: done
pop from [1 3 2 7 4 5]
  took 1, moved 5 to root -> [5 3 2 7 4]
  swap with child at 2 -> [2 3 5 7 4]
  no smaller child: done
pop from [2 3 5 7 4]
  took 2, moved 4 to root -> [4 3 5 7]
  swap with child at 1 -> [3 4 5 7]
  no smaller child: done

The popped values come out in order: 0, then 1, then 2. Popping everything from a heap is how heap sort works.

Step through it

InteractiveMin-heap: push and popPush appends and sifts up; pop moves the last element to the root and sifts down. The tree is stored in a plain slice.

slice

Push 5: append it at index 0, the next free spot on the bottom level.

1 / 26

container/heap#

Go gives you the sift-up and sift-down algorithms in container/heap. You supply a type with five methods: sort.Interface (Len, Less, Swap) plus Push and Pop that just append and remove the last element. The package does the reordering.

type MinHeap []int
 
func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] } // flip for a max-heap
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

Then always go through the package functions: heap.Init(h), heap.Push(h, x), heap.Pop(h). The minimum is (*h)[0].

Try it: a priority queue of structs

Less decides what "highest priority" means. Here a larger Priority wins, so Less uses >, which makes it a max-heap:

pq.go

package main
 
import (
	"container/heap"
	"fmt"
)
 
type Task struct {
	Name     string
	Priority int
}
 
type PQ []Task
 
func (q PQ) Len() int           { return len(q) }
func (q PQ) Less(i, j int) bool { return q[i].Priority > q[j].Priority }
func (q PQ) Swap(i, j int)      { q[i], q[j] = q[j], q[i] }
func (q *PQ) Push(x any)        { *q = append(*q, x.(Task)) }
func (q *PQ) Pop() any {
	old := *q
	x := old[len(old)-1]
	*q = old[:len(old)-1]
	return x
}
 
func main() {
	q := &PQ{{"reply email", 1}, {"fix outage", 9}, {"write tests", 5}}
	heap.Init(q) // O(n): arrange existing items into a heap
	heap.Push(q, Task{"review PR", 7})
	fmt.Println("next up:", (*q)[0].Name)
	for q.Len() > 0 {
		t := heap.Pop(q).(Task)
		fmt.Printf("%d %s\n", t.Priority, t.Name)
	}
}

Output

next up: fix outage
9 fix outage
7 review PR
5 write tests
1 reply email

Top-k: keep a small heap#

"Find the k-th largest element." Sorting costs O(n log n). Instead keep a min-heap of size k holding the k largest seen so far. For each number, push it; if the heap grows past k, pop the smallest. At the end, the root is the k-th largest. Cost: O(n log k), and only O(k) memory, which works even on a stream too big to hold.

Trace [3, 2, 1, 5, 6, 4], k = 2:

push	heap after push	over k? pop	heap (2 largest so far)
3	[3]		[3]
2	[2 3]		[2 3]
1	[1 3 2]	pop 1	[2 3]
5	[2 3 5]	pop 2	[3 5]
6	[3 5 6]	pop 3	[5 6]
4	[4 6 5]	pop 4	[5 6]

Root is 5: the 2nd largest.

Run the trace

topk_trace.go

package main
 
import (
	"container/heap"
	"fmt"
	"slices"
)
 
type MinHeap []int
 
func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}
 
func main() {
	nums, k := []int{3, 2, 1, 5, 6, 4}, 2
	h := &MinHeap{}
	for _, n := range nums {
		heap.Push(h, n)
		line := fmt.Sprintf("push %d", n)
		if h.Len() > k {
			line += fmt.Sprintf(", pop %d", heap.Pop(h))
		}
		kept := slices.Clone(*h)
		slices.Sort(kept)
		fmt.Printf("%-14s %d largest so far: %v\n", line, k, kept)
	}
	fmt.Println("k-th largest:", (*h)[0])
}

Output

push 3         2 largest so far: [3]
push 2         2 largest so far: [2 3]
push 1, pop 1  2 largest so far: [2 3]
push 5, pop 2  2 largest so far: [3 5]
push 6, pop 3  2 largest so far: [5 6]
push 4, pop 4  2 largest so far: [5 6]
k-th largest: 5

Merging sorted streams#

To merge k sorted lists, keep a heap of the current front of each list. Pop the smallest, output it, push the next element from the same list. Each element passes through the heap once: O(N log k) for N total elements.

Try it: merge k sorted lists

The heap holds one entry per list: its current front value, plus which list and position it came from:

mergek.go

package main
 
import (
	"container/heap"
	"fmt"
)
 
type item struct{ val, list, idx int }
 
type H []item
 
func (h H) Len() int           { return len(h) }
func (h H) Less(i, j int) bool { return h[i].val < h[j].val }
func (h H) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *H) Push(x any)        { *h = append(*h, x.(item)) }
func (h *H) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}
 
func main() {
	lists := [][]int{{1, 4, 7}, {2, 5, 8}, {0, 6}}
	h := &H{}
	for li, l := range lists {
		heap.Push(h, item{l[0], li, 0})
	}
	var out []int
	for h.Len() > 0 {
		it := heap.Pop(h).(item)
		out = append(out, it.val)
		if next := it.idx + 1; next < len(lists[it.list]) {
			heap.Push(h, item{lists[it.list][next], it.list, next})
		}
		fmt.Printf("took %d from list %d, heap size %d\n", it.val, it.list, h.Len())
	}
	fmt.Println(out)
}

Output

took 0 from list 2, heap size 3
took 1 from list 0, heap size 3
took 2 from list 1, heap size 3
took 4 from list 0, heap size 3
took 5 from list 1, heap size 3
took 6 from list 2, heap size 2
took 7 from list 0, heap size 1
took 8 from list 1, heap size 0
[0 1 2 4 5 6 7 8]

Heaps over streams#

Heaps shine when data arrives over time and you cannot sort it once and be done.

K-th largest element in a stream

Keep a min-heap holding the k largest values seen so far. Its root is the k-th largest. Each new value is pushed; if the heap grows beyond k, pop the smallest. Every Add is O(log k), whatever the stream's length.

Median of a data stream: two heaps

Split the values into a lower half and an upper half:

low is a max-heap (its root is the largest of the small values);
high is a min-heap (its root is the smallest of the large values);
keep len(low) == len(high) or len(low) == len(high) + 1.

The median is then low's root (odd count) or the average of both roots (even count). Adding a value: push it into low if it is at most low's root, otherwise into high, then move one root across if the sizes drift apart.

InteractiveMedian of a data stream (two heaps)A max-heap holds the smaller half and a min-heap the larger half. Keep them balanced and the median sits on top.

median 5seen 1

stream

low (max-heap)

high (min-heap)

empty

Add 5 to low (it is ≤ low's max). Median = top of low = 5.

1 / 9

Top K frequent elements

Count with a map, then keep a min-heap of size k ordered by count: O(n log k). The hashing lesson showed the O(n) bucket-sort version; the heap version also works when counts arrive as a stream.

Merge K sorted linked lists

The same heap-of-fronts idea from the slices example above, with list nodes: push each list's head; pop the smallest node, append it to the result, and push its Next.

streams.go

package main
 
import (
	"container/heap"
	"fmt"
	"sort"
)
 
// IntHeap is a min-heap; wrap values with a minus sign to use it as a max-heap.
type IntHeap []int
 
func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *IntHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}
 
// --- K-th largest in a stream ---
 
type KthLargest struct {
	k int
	h IntHeap
}
 
func (s *KthLargest) Add(v int) int {
	heap.Push(&s.h, v)
	if s.h.Len() > s.k {
		heap.Pop(&s.h)
	}
	return s.h[0]
}
 
// --- Running median ---
 
type Median struct{ low, high IntHeap } // low stores negated values: a max-heap
 
func (m *Median) Add(v int) {
	if m.low.Len() == 0 || v <= -m.low[0] {
		heap.Push(&m.low, -v)
	} else {
		heap.Push(&m.high, v)
	}
	if m.low.Len() > m.high.Len()+1 {
		heap.Push(&m.high, -heap.Pop(&m.low).(int))
	} else if m.high.Len() > m.low.Len() {
		heap.Push(&m.low, -heap.Pop(&m.high).(int))
	}
}
 
func (m *Median) Value() float64 {
	if m.low.Len() > m.high.Len() {
		return float64(-m.low[0])
	}
	return float64(-m.low[0]+m.high[0]) / 2
}
 
// --- Top K frequent with a size-k heap ---
 
type pair struct{ val, count int }
type pairHeap []pair
 
func (h pairHeap) Len() int           { return len(h) }
func (h pairHeap) Less(i, j int) bool { return h[i].count < h[j].count }
func (h pairHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *pairHeap) Push(x any)        { *h = append(*h, x.(pair)) }
func (h *pairHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}
 
func topK(nums []int, k int) []int {
	count := map[int]int{}
	for _, n := range nums {
		count[n]++
	}
	h := &pairHeap{}
	for v, c := range count {
		heap.Push(h, pair{v, c})
		if h.Len() > k {
			heap.Pop(h) // drop the least frequent
		}
	}
	out := make([]int, 0, k)
	for _, p := range *h {
		out = append(out, p.val)
	}
	sort.Ints(out) // map order is random; sort for stable output
	return out
}
 
// --- Merge K sorted linked lists ---
 
type ListNode struct {
	Val  int
	Next *ListNode
}
 
type nodeHeap []*ListNode
 
func (h nodeHeap) Len() int           { return len(h) }
func (h nodeHeap) Less(i, j int) bool { return h[i].Val < h[j].Val }
func (h nodeHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *nodeHeap) Push(x any)        { *h = append(*h, x.(*ListNode)) }
func (h *nodeHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}
 
func mergeKLists(lists []*ListNode) *ListNode {
	h := &nodeHeap{}
	for _, l := range lists {
		if l != nil {
			heap.Push(h, l)
		}
	}
	dummy := &ListNode{}
	tail := dummy
	for h.Len() > 0 {
		n := heap.Pop(h).(*ListNode)
		tail.Next = n
		tail = n
		if n.Next != nil {
			heap.Push(h, n.Next)
		}
	}
	return dummy.Next
}
 
func list(vals ...int) *ListNode {
	dummy := &ListNode{}
	t := dummy
	for _, v := range vals {
		t.Next = &ListNode{Val: v}
		t = t.Next
	}
	return dummy.Next
}
 
func main() {
	kth := &KthLargest{k: 3}
	for _, v := range []int{4, 5, 8, 2} {
		kth.Add(v)
	}
	fmt.Print("3rd largest after each add:")
	for _, v := range []int{3, 5, 10, 9, 4} {
		fmt.Print(" ", kth.Add(v))
	}
	fmt.Println()
 
	var m Median
	fmt.Print("running median:")
	for _, v := range []int{5, 15, 1, 3, 8, 7} {
		m.Add(v)
		fmt.Print(" ", m.Value())
	}
	fmt.Println()
 
	fmt.Println("top 2 frequent:", topK([]int{1, 1, 1, 2, 2, 3, 4, 4, 4, 4}, 2))
 
	merged := mergeKLists([]*ListNode{list(1, 4, 5), list(1, 3, 4), list(2, 6)})
	fmt.Print("merged lists:")
	for n := merged; n != nil; n = n.Next {
		fmt.Print(" ", n.Val)
	}
	fmt.Println()
}

Output

3rd largest after each add: 4 5 5 8 8
running median: 5 10 5 4 5 6
top 2 frequent: [1 4]
merged lists: 1 1 2 3 4 4 5 6

Median stores the lower half negated so the same min-heap type acts as a max-heap. That trick saves writing a second heap type, but it fails for math.MinInt, whose negation overflows; in that case write a separate type with > in Less.

The full program#

heaps.go

package main
 
import (
	"container/heap"
	"fmt"
)
 
type MinHeap []int
 
func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}
 
func findKthLargest(nums []int, k int) int {
	h := &MinHeap{}
	for _, n := range nums {
		heap.Push(h, n)
		if h.Len() > k {
			heap.Pop(h)
		}
	}
	return (*h)[0]
}
 
// A generic heap, the kind you will put in your own library.
type Heap[T any] struct {
	data []T
	less func(a, b T) bool
}
 
func (h *Heap[T]) Len() int { return len(h.data) }
 
func (h *Heap[T]) Push(v T) {
	h.data = append(h.data, v)
	i := len(h.data) - 1
	for i > 0 {
		p := (i - 1) / 2
		if !h.less(h.data[i], h.data[p]) {
			break
		}
		h.data[i], h.data[p] = h.data[p], h.data[i]
		i = p
	}
}
 
func (h *Heap[T]) Pop() T {
	top := h.data[0]
	last := len(h.data) - 1
	h.data[0] = h.data[last]
	h.data = h.data[:last]
	i := 0
	for {
		l, r, m := 2*i+1, 2*i+2, i
		if l < len(h.data) && h.less(h.data[l], h.data[m]) {
			m = l
		}
		if r < len(h.data) && h.less(h.data[r], h.data[m]) {
			m = r
		}
		if m == i {
			return top
		}
		h.data[i], h.data[m] = h.data[m], h.data[i]
		i = m
	}
}
 
type task struct {
	name     string
	priority int
}
 
func lastStoneWeight(stones []int) int {
	h := &Heap[int]{less: func(a, b int) bool { return a > b }} // max-heap
	for _, s := range stones {
		h.Push(s)
	}
	for h.Len() > 1 {
		a, b := h.Pop(), h.Pop()
		if a != b {
			h.Push(a - b)
		}
	}
	if h.Len() == 0 {
		return 0
	}
	return h.Pop()
}
 
func main() {
	h := &MinHeap{5, 3, 7, 1}
	heap.Init(h) // reorder an existing slice in O(n)
	heap.Push(h, 0)
	fmt.Print("min-heap pops: ")
	for h.Len() > 0 {
		fmt.Print(heap.Pop(h), " ")
	}
	fmt.Println()
 
	fmt.Println(findKthLargest([]int{3, 2, 1, 5, 6, 4}, 2))
 
	pq := &Heap[task]{less: func(a, b task) bool { return a.priority > b.priority }}
	pq.Push(task{"reply email", 1})
	pq.Push(task{"fix outage", 5})
	pq.Push(task{"write tests", 3})
	for pq.Len() > 0 {
		fmt.Print(pq.Pop().name, "; ")
	}
	fmt.Println()
 
	fmt.Println(lastStoneWeight([]int{2, 7, 4, 1, 8, 1}))
}

Output

min-heap pops: 0 1 3 5 7 
5
fix outage; write tests; reply email; 
1

Complexity#

Operation	Time
peek min (`h[0]`)	O(1)
push, pop	O(log n)
build a heap from n items (`heap.Init`)	O(n)
top-k from n items	O(n log k)
merge k sorted lists of N total items	O(N log k)

Recognizing the pattern#

"k largest / smallest / closest / most frequent", "repeatedly take the best remaining", "merge sorted inputs", "schedule by priority" and "running median" all mean a heap. Dijkstra's shortest paths (two lessons from now) is a heap at its core.

Summary#

A heap keeps every parent at most its children, so the min is at the root.
Stored in a slice: children at 2i+1 and 2i+2, parent at (i-1)/2.
Push sifts up, pop sifts down: O(log n) each.
container/heap needs Len, Less, Swap, Push, Pop; use the package functions to modify it.
Top-k with a size-k heap is O(n log k) and works on streams.

Heaps & priority queues

The heap property#

Push: add at the end, sift up#

Run the trace

Pop: take the root, sift down#

Run the trace

Step through it

container/heap#

Try it: a priority queue of structs

Top-k: keep a small heap#

Run the trace

Merging sorted streams#

Try it: merge k sorted lists

Heaps over streams#

K-th largest element in a stream

Median of a data stream: two heaps

Top K frequent elements

Merge K sorted linked lists

The full program#

Complexity#

Recognizing the pattern#

Summary#

Check your understanding

Practice in Go