Heaps & Priority Queues

Code

The harness and the tests

Three files: the harness, the tests, and two new commands.

command	what it does
`algo top`	the best k: heap against "sort and take"
`algo pq`	lesson 6's question: your heap against `container/heap`

heapbench.go also carries stdItems — the interface container/heap demands. Look at its signatures now; in step 6 they are the entire answer:

func (s *stdItems) Push(x any)  { *s = append(*s, x.(Item)) }
func (s *stdItems) Pop() any    { ... }

The tests

test	what it requires
`TestPopComesOutAscending`	the invariant holds after every `Push`; draining gives ascending order
`TestHeightIsImpliedByLength`	`Height()` is `floor(log2 n)+1` for every n up to 5,000
`TestTheHeapIsNotSorted`	the array is NOT sorted — and is still a heap
`TestTopKAgreesWithSortingAndCostsLess`	the same k ratings, in under `2n` comparisons
`TestTopKLosesWhenKApproachesN`	at `k = n` the heap loses to sorting

Two of them demand that something is not the case, and both protect the point of the lesson.

"Not sorted" guards against an improvement that would ruin everything. Turn Push into an insertion sort and every other test still passes, while the cost goes from O(log n) to O(n). A heap does not sort — that is not a shortcoming, it is exactly what it sells in exchange for the speed.

"Loses at k = n" is the honest boundary. The heap wins only while k is well under n; step 5 measures where it turns over. The test stops that being quietly left out.

The test file also carries four Benchmarks — step 6 needs them, because what it measures is something the counter cannot see:

go test -run XXX -bench 'Heap|TopK' -benchtime=20x -count=5

heapbench.go CREATE

Create this file in ./

package main

import (
	stdheap "container/heap"
	"flag"
	"fmt"
	"sort"
	"time"

	"algo/metrics"
)

// stdItems is the adapter container/heap requires. Note the shape of Push and
// Pop: they take and return `any`, so every element crosses an interface
// boundary on the way in and out.
type stdItems []Item

func (s stdItems) Len() int           { return len(s) }
func (s stdItems) Less(i, j int) bool { return s[i].Rating < s[j].Rating }
func (s stdItems) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s *stdItems) Push(x any)        { *s = append(*s, x.(Item)) }
func (s *stdItems) Pop() any          { old := *s; n := len(old); it := old[n-1]; *s = old[:n-1]; return it }

// topKStd is TopKHeap written against container/heap. Same algorithm, same
// comparisons — a different way of reaching the elements.
func topKStd(items []Item, k int) []Item {
	if k <= 0 {
		return nil
	}
	h := &stdItems{}
	stdheap.Init(h)
	for _, it := range items {
		if h.Len() < k {
			stdheap.Push(h, it)
			continue
		}
		if it.Rating <= (*h)[0].Rating {
			continue
		}
		stdheap.Pop(h)
		stdheap.Push(h, it)
	}
	out := make([]Item, 0, h.Len())
	for h.Len() > 0 {
		out = append(out, stdheap.Pop(h).(Item))
	}
	for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
		out[i], out[j] = out[j], out[i]
	}
	return out
}

// topKSort is the obvious answer: sort everything, take the first k.
func topKSort(items []Item, k int, c *metrics.Counter) []Item {
	cp := append([]Item(nil), items...)
	sort.Slice(cp, func(i, j int) bool {
		c.Hit()
		return cp[i].Rating > cp[j].Rating
	})
	if k > len(cp) {
		k = len(cp)
	}
	return cp[:k]
}

// cmdTop is the partial-order argument: you do not have to sort to find the best
// ten.
func cmdTop(args []string) error {
	fs := flag.NewFlagSet("top", flag.ContinueOnError)
	in := fs.String("in", "library.jsonl", "library file")
	k := fs.Int("k", 10, "how many to take")
	show := fs.Bool("show", false, "print the k items")
	if err := fs.Parse(args); err != nil {
		return err
	}
	items, err := load(*in)
	if err != nil {
		return err
	}
	n := len(items)

	var hc metrics.Counter
	t0 := time.Now()
	byHeap := TopKHeap(items, *k, &hc)
	heapTime := time.Since(t0)

	var sc metrics.Counter
	t0 = time.Now()
	bySort := topKSort(items, *k, &sc)
	sortTime := time.Since(t0)

	fmt.Printf("n = %d, k = %d\n\n", n, *k)
	fmt.Printf("  %-22s %14s %12s %14s\n", "", "comparisons", "time", "per item")
	fmt.Printf("  %-22s %14d %12s %14.2f\n", "heap of size k", hc.N(), heapTime.Round(time.Microsecond), float64(hc.N())/float64(n))
	fmt.Printf("  %-22s %14d %12s %14.2f\n", "sort.Slice, take k", sc.N(), sortTime.Round(time.Microsecond), float64(sc.N())/float64(n))
	fmt.Printf("\n  ratio: %.1fx fewer comparisons\n", float64(sc.N())/float64(hc.N()))

	same := len(byHeap) == len(bySort)
	for i := range byHeap {
		if i < len(bySort) && byHeap[i].Rating != bySort[i].Rating {
			same = false
		}
	}
	fmt.Printf("  same k ratings? %v\n", same)
	if *show {
		fmt.Println()
		for i, it := range byHeap {
			fmt.Printf("  %2d. %3d  %s\n", i+1, it.Rating, it.Title)
		}
	}
	return nil
}

// cmdPq is lesson 6's unanswered question: the hand-written heap against
// container/heap, on the workload a priority queue actually gets.
func cmdPq(args []string) error {
	fs := flag.NewFlagSet("pq", flag.ContinueOnError)
	in := fs.String("in", "library.jsonl", "library file")
	rounds := fs.Int("rounds", 20, "how many times to repeat each arm")
	if err := fs.Parse(args); err != nil {
		return err
	}
	items, err := load(*in)
	if err != nil {
		return err
	}

	// push everything, then drain — the shape of Dijkstra's frontier.
	var mine, std time.Duration
	var c metrics.Counter
	for r := 0; r < *rounds; r++ {
		t0 := time.Now()
		h := NewMinHeap(len(items))
		for _, it := range items {
			h.Push(it, &c)
		}
		for h.Len() > 0 {
			h.Pop(&c)
		}
		mine += time.Since(t0)

		t0 = time.Now()
		s := &stdItems{}
		stdheap.Init(s)
		for _, it := range items {
			stdheap.Push(s, it)
		}
		for s.Len() > 0 {
			stdheap.Pop(s)
		}
		std += time.Since(t0)
	}
	n := time.Duration(*rounds)
	fmt.Printf("n = %d, push-all then drain, %d rounds\n\n", len(items), *rounds)
	fmt.Printf("  %-22s %14s\n", "", "per round")
	fmt.Printf("  %-22s %14s\n", "hand-written heap", (mine / n).Round(time.Microsecond))
	fmt.Printf("  %-22s %14s\n", "container/heap", (std / n).Round(time.Microsecond))
	fmt.Printf("\n  container/heap is %.2fx slower\n", float64(std)/float64(mine))
	return nil
}

heap_test.go CREATE

Create this file in ./

package main

import (
	stdheap "container/heap"
	"math"
	"sort"
	"testing"

	"algo/metrics"
)

// scatteredRatings is a deterministic xorshift rating set. Lesson 9's ratedLib
// walks the ratings in a cycle, which is fine for a counting sort and useless
// here — a heap fed an orderly sequence never has to move anything.
func scatteredRatings(n int) []Item {
	s := make([]Item, n)
	x := uint64(88172645463325252)
	for i := range s {
		x ^= x << 13
		x ^= x >> 7
		x ^= x << 17
		s[i] = Item{ID: i + 1, Rating: int(x % (MaxRating + 1))}
	}
	return s
}

// A heap is not sorted, but draining it IS.
func TestPopComesOutAscending(t *testing.T) {
	items := scatteredRatings(5000)
	h := NewMinHeap(0)
	var c metrics.Counter
	for _, it := range items {
		h.Push(it, &c)
		if !h.IsHeap() {
			t.Fatalf("the heap invariant broke during Push at len %d", h.Len())
		}
	}
	if h.Len() != len(items) {
		t.Fatalf("len %d, want %d", h.Len(), len(items))
	}
	prev := -1
	for h.Len() > 0 {
		it, _ := h.Pop(&c)
		if it.Rating < prev {
			t.Fatalf("Pop returned %d after %d — not ascending", it.Rating, prev)
		}
		prev = it.Rating
	}
}

// THE ARRAY IS THE TREE. The shape is implied by the length alone: no traversal
// is needed to know the height, because there is nothing to traverse.
func TestHeightIsImpliedByLength(t *testing.T) {
	h := NewMinHeap(0)
	var c metrics.Counter
	for n := 1; n <= 5000; n++ {
		h.Push(Item{Rating: n % (MaxRating + 1)}, &c)
		want := int(math.Floor(math.Log2(float64(n)))) + 1
		if h.Height() != want {
			t.Fatalf("n=%d: Height %d, want %d", n, h.Height(), want)
		}
	}
}

// A HEAP IS NOT SORTED, AND THAT IS THE POINT. It knows its minimum and nothing
// else — asserted here so nobody "improves" Push into an insertion sort.
func TestTheHeapIsNotSorted(t *testing.T) {
	items := scatteredRatings(1000)
	h := NewMinHeap(0)
	var c metrics.Counter
	for _, it := range items {
		h.Push(it, &c)
	}
	got := make([]int, 0, h.Len())
	for _, it := range h.a {
		got = append(got, it.Rating)
	}
	if sort.IntsAreSorted(got) {
		t.Fatal("the array came out fully sorted — that is a stronger property " +
			"than a heap has, and it costs more than a heap costs")
	}
	// It is a heap, though: every parent is <= its children.
	if !h.IsHeap() {
		t.Fatal("the heap invariant does not hold")
	}
	min, _ := h.Peek()
	for _, it := range items {
		if it.Rating < min.Rating {
			t.Fatalf("Peek returned %d but %d is present", min.Rating, it.Rating)
		}
	}
}

// TopKHeap must agree with sorting, and cost far less.
func TestTopKAgreesWithSortingAndCostsLess(t *testing.T) {
	items := scatteredRatings(100000)
	const k = 10

	var hc metrics.Counter
	byHeap := TopKHeap(items, k, &hc)
	var sc metrics.Counter
	bySort := topKSort(items, k, &sc)

	if len(byHeap) != k {
		t.Fatalf("heap returned %d items, want %d", len(byHeap), k)
	}
	for i := range byHeap {
		if byHeap[i].Rating != bySort[i].Rating {
			t.Fatalf("rank %d: heap %d, sort %d", i, byHeap[i].Rating, bySort[i].Rating)
		}
	}
	// O(n log k) against O(n log n): with k=10 and n=100000 the heap should be
	// within a small multiple of n, while the sort is not.
	if float64(hc.N()) > 2*float64(len(items)) {
		t.Errorf("%d comparisons for k=%d, want about n=%d", hc.N(), k, len(items))
	}
	if hc.N() >= sc.N() {
		t.Errorf("heap %d comparisons, sort %d — the heap should win by a lot", hc.N(), sc.N())
	}
	t.Logf("n=%d k=%d: heap %d comparisons, sort %d (%.1fx)",
		len(items), k, hc.N(), sc.N(), float64(sc.N())/float64(hc.N()))
}

// AND THE HONEST BOUNDARY: at large k the heap LOSES. O(n log k) is only a win
// while log k is well under log n.
func TestTopKLosesWhenKApproachesN(t *testing.T) {
	items := scatteredRatings(100000)
	var hc metrics.Counter
	TopKHeap(items, len(items), &hc)
	var sc metrics.Counter
	topKSort(items, len(items), &sc)
	if hc.N() <= sc.N() {
		t.Errorf("at k=n the heap used %d comparisons and the sort %d — the heap "+
			"is expected to lose here", hc.N(), sc.N())
	}
	t.Logf("k=n: heap %d comparisons, sort %d (%.1fx worse)",
		hc.N(), sc.N(), float64(hc.N())/float64(sc.N()))
}

// LESSON 6'S UNANSWERED QUESTION. There the standard library won; here it does
// not, and the benchmarks are where that shows.
//
//	go test -run XXX -bench 'Heap' -benchtime=20x -count=5
func benchHeapItems() []Item { return scatteredRatings(10000) }

func BenchmarkHeapMine(b *testing.B) {
	items := benchHeapItems()
	var c metrics.Counter
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		h := NewMinHeap(len(items))
		for _, it := range items {
			h.Push(it, &c)
		}
		for h.Len() > 0 {
			h.Pop(&c)
		}
	}
}

func BenchmarkHeapStdlib(b *testing.B) {
	items := benchHeapItems()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s := &stdItems{}
		stdheap.Init(s)
		for _, it := range items {
			stdheap.Push(s, it)
		}
		for s.Len() > 0 {
			stdheap.Pop(s)
		}
	}
}

func BenchmarkTopKHeap(b *testing.B) {
	items := scatteredRatings(100000)
	var c metrics.Counter
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		TopKHeap(items, 10, &c)
	}
}

func BenchmarkTopKSort(b *testing.B) {
	items := scatteredRatings(100000)
	var c metrics.Counter
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		topKSort(items, 10, &c)
	}
}

main.go MODIFY 6 lines changed

Edit this file in ./

package main

import (
	"fmt"
	"os"
)

func usage() {
	fmt.Fprint(os.Stderr, `algo — your library workbench

  algo gen  [-n N] [-seed S] [-o FILE] [-sorted|-reverse|-shuffled] [-dup-rate R]
  algo find -title TITLE [-in FILE]
  algo load [-in FILE]
  algo stats -growth [-n N]
  algo playlist -from I -to J [-in FILE]
  algo history [-n N] [-in FILE]
  algo queue [-n N]
  algo walk [-depth D] [-branch B] [-per P]
  algo depth -n N [-iterative]
  algo sweep [-scattered]
  algo std [-n N]
  algo sort [-algo A] [-order O] [-in FILE]
  algo dc [-order O] [-in FILE]
  algo stable [-n N]
  algo count [-in FILE]
  algo radix [-in FILE]
  algo index [-in FILE] [-nogrow] [-buckets N]
  algo ordered [-in FILE] [-n N]
  algo alpha [-in FILE] [-n N] [-from A -to B]
  algo tree [-in FILE]
  algo shape [-shuffled FILE] [-sorted FILE]
  algo balance [-shuffled FILE] [-sorted FILE]
  algo rotate [-in FILE]
  algo reload [-in FILE]
  algo top [-in FILE] [-k N] [-show]
  algo pq [-in FILE] [-rounds N]
`)
}

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(2)
	}
	var err error
	switch os.Args[1] {
	case "gen":
		err = cmdGen(os.Args[2:])
	case "find":
		err = cmdFind(os.Args[2:])
	case "load":
		err = cmdLoad(os.Args[2:])
	case "stats":
		err = cmdStats(os.Args[2:])
	case "playlist":
		err = cmdPlaylist(os.Args[2:])
	case "history":
		err = cmdHistory(os.Args[2:])
	case "queue":
		err = cmdQueue(os.Args[2:])
	case "walk":
		err = cmdWalk(os.Args[2:])
	case "depth":
		err = cmdDepth(os.Args[2:])
	case "sweep":
		err = cmdSweep(os.Args[2:])
	case "std":
		err = cmdStd(os.Args[2:])
	case "sort":
		err = cmdSort(os.Args[2:])
	case "dc":
		err = cmdDC(os.Args[2:])
	case "stable":
		err = cmdStable(os.Args[2:])
	case "count":
		err = cmdCount(os.Args[2:])
	case "radix":
		err = cmdRadix(os.Args[2:])
	case "index":
		err = cmdIndex(os.Args[2:])
	case "ordered":
		err = cmdOrdered(os.Args[2:])
	case "alpha":
		err = cmdAlpha(os.Args[2:])
	case "tree":
		err = cmdTree(os.Args[2:])
	case "shape":
		err = cmdShape(os.Args[2:])
	case "balance":
		err = cmdBalance(os.Args[2:])
	case "rotate":
		err = cmdRotate(os.Args[2:])
	case "reload":
		err = cmdReload(os.Args[2:])
	case "top":
		err = cmdTop(os.Args[2:])
	case "pq":
		err = cmdPq(os.Args[2:])
	default:
		usage()
		os.Exit(2)
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, "algo:", err)
		os.Exit(1)
	}
}