aboutsummaryrefslogtreecommitdiffstats
path: root/heap
diff options
context:
space:
mode:
Diffstat (limited to 'heap')
-rw-r--r--heap/example_test.go178
-rw-r--r--heap/heap.go87
-rw-r--r--heap/heap_test.go72
3 files changed, 337 insertions, 0 deletions
diff --git a/heap/example_test.go b/heap/example_test.go
new file mode 100644
index 0000000..f8ea7d8
--- /dev/null
+++ b/heap/example_test.go
@@ -0,0 +1,178 @@
+package heap_test
+
+import (
+ "cmp"
+ "fmt"
+
+ "github.com/rhogenson/container/heap"
+)
+
+func Example_dijkstra() {
+ const (
+ maze = `
+---------------------
+ | | |
+| --- | | --- ----- |
+| | | | | |
+|-- |-----| ----- --|
+| | | | | |
+| --- --- | --- | | |
+| | | | | | | |
+| | --- | | --- | | |
+| | | | | | | |
+| | --- | |-- --|-- |
+| | | | | | | |
+| |-- | | | --- | --|
+| | | | | | |
+| | --| ----- ----- |
+| | | | |
+| | | | --------- | |
+| | | | | | | |
+| --- |---- | | | | |
+| | |
+---------------------
+`
+ width = 21
+ height = 21
+ )
+
+ type point struct{ x, y int }
+ start := point{0, 1}
+ goal := point{20, 19}
+ neighbors := func(p point) []point {
+ return []point{
+ {p.x - 1, p.y},
+ {p.x + 1, p.y},
+ {p.x, p.y - 1},
+ {p.x, p.y + 1},
+ }
+ }
+ walkable := func(p point) bool {
+ return 0 <= p.y && p.y < height &&
+ 0 <= p.x && p.x < width &&
+ maze[p.y*(width+1)+p.x+1] == ' '
+ }
+
+ visited := map[point]int{start: 1}
+ q := heap.New(func(x, y point) int { return cmp.Compare(visited[x], visited[y]) })
+ q.Push(start)
+Dijkstra:
+ for {
+ p, ok := q.Pop()
+ if !ok {
+ fmt.Println("Giving up!")
+ return
+ }
+ for _, neighbor := range neighbors(p) {
+ if !walkable(neighbor) || visited[neighbor] > 0 {
+ continue
+ }
+ visited[neighbor] = visited[p] + 1
+ if neighbor == goal {
+ break Dijkstra
+ }
+ q.Push(neighbor)
+ }
+ }
+
+ completedMaze := []byte(maze)
+ fillIn := func(p point) {
+ completedMaze[p.y*(width+1)+p.x+1] = '*'
+ }
+ for p := goal; p != start; {
+ fillIn(p)
+ closestPoint := p
+ for _, neighbor := range neighbors(p) {
+ if walkable(neighbor) && visited[neighbor] > 0 && visited[neighbor] < visited[closestPoint] {
+ closestPoint = neighbor
+ }
+ }
+ p = closestPoint
+ }
+ fillIn(start)
+
+ fmt.Printf("%s\n", completedMaze)
+
+ // Output:
+ // ---------------------
+ // ** | | *******|
+ // |*--- | | ---*-----*|
+ // |***| | |***| ***|
+ // |--*|-----|*-----*--|
+ // |***|*****|* |*| |
+ // |*---*---*|*--- |*| |
+ // |*|***| |*** | |*| |
+ // |*|*--- | | --- |*| |
+ // |*|*| | | | |***|
+ // |*|*--- | |-- --|--*|
+ // |*|***| | | | |***|
+ // |*|--*| | | --- |*--|
+ // |*|***| | |*****| |
+ // |*|*--| -----*----- |
+ // |*|***|*******| |
+ // |*| |*|*--------- | |
+ // |*| |*|*****|***| | |
+ // |*---*|----*|*|*| | |
+ // |*****| ***|******
+ // ---------------------
+}
+
+func ExampleNew() {
+ priority := map[string]int{
+ "job1": 10,
+ "job2": 30,
+ "job3": 100,
+ "job4": 20,
+ }
+ h := heap.New(func(j1, j2 string) int { return cmp.Compare(priority[j1], priority[j2]) })
+ h.Push("job1")
+ h.Push("job2")
+ h.Push("job3")
+ h.Push("job4")
+ if highestPriorityJob, ok := h.Pop(); ok {
+ fmt.Println(highestPriorityJob)
+ }
+
+ // Output:
+ // job1
+}
+
+func ExampleHeap_Len() {
+ h := heap.New(cmp.Compare[int])
+ h.Push(1)
+ h.Push(2)
+ h.Push(3)
+ fmt.Println(h.Len())
+
+ // Output:
+ // 3
+}
+
+func ExampleHeap_Grow() {
+ h := heap.New(cmp.Compare[int])
+ h.Grow(3)
+ // Push without allocating:
+ h.Push(1)
+ h.Push(2)
+ h.Push(3)
+}
+
+func ExampleHeap_Push() {
+ h := heap.New(cmp.Compare[int])
+ h.Push(1)
+ h.Push(2)
+ h.Push(3)
+}
+
+func ExampleHeap_Pop() {
+ h := heap.New(cmp.Compare[int])
+ h.Push(1)
+ h.Push(2)
+ h.Push(3)
+ if n, ok := h.Pop(); ok {
+ fmt.Println(n)
+ }
+
+ // Output:
+ // 1
+}
diff --git a/heap/heap.go b/heap/heap.go
new file mode 100644
index 0000000..612d5a8
--- /dev/null
+++ b/heap/heap.go
@@ -0,0 +1,87 @@
+// Package heap implements a priority queue as a min heap backed by a slice.
+//
+// This can be seen as a replacement for the standard library [container/heap]
+// package which was created before generics were a thing.
+package heap
+
+import (
+ "slices"
+)
+
+// Heap is a binary heap backed by a slice.
+type Heap[T any] struct {
+ buf []T
+ compare func(T, T) int
+}
+
+// New creates a new heap with the given comparison function.
+func New[T any](compare func(T, T) int) *Heap[T] {
+ return &Heap[T]{compare: compare}
+}
+
+// Len returns the number of elements in the Heap.
+func (h *Heap[T]) Len() int {
+ return len(h.buf)
+}
+
+// Grow makes space for at least n more elements to be pushed onto the heap
+// without reallocating.
+func (h *Heap[T]) Grow(n int) {
+ h.buf = slices.Grow(h.buf, n)
+}
+
+// Push pushes the element x onto the heap.
+// The complexity is O(log n) where n = h.Len().
+func (h *Heap[T]) Push(x T) {
+ n := len(h.buf)
+ h.buf = append(h.buf, x)
+ h.up(n)
+}
+
+// Pop removes and returns the minimum element (according to Less) from
+// the heap. The complexity is O(log n) where n = h.Len().
+func (h *Heap[T]) Pop() (T, bool) {
+ if len(h.buf) == 0 {
+ var zero T
+ return zero, false
+ }
+ x := h.buf[0]
+ last := h.buf[len(h.buf)-1]
+ h.buf = h.buf[:len(h.buf)-1]
+ if len(h.buf) > 0 {
+ h.down(0, last)
+ }
+ return x, true
+}
+
+func (h *Heap[T]) up(j int) {
+ x := h.buf[j]
+ for {
+ i := (j - 1) / 2 // parent
+ if i == j || h.compare(x, h.buf[i]) >= 0 {
+ break
+ }
+ h.buf[j] = h.buf[i]
+ j = i
+ }
+ h.buf[j] = x
+}
+
+func (h *Heap[T]) down(i int, x T) {
+ for {
+ j1 := 2*i + 1
+ if j1 >= len(h.buf) || j1 < 0 { // j1 < 0 after int overflow
+ break
+ }
+ j := j1 // left child
+ if j2 := j1 + 1; j2 < len(h.buf) && h.compare(h.buf[j2], h.buf[j1]) < 0 {
+ j = j2 // = 2*i + 2 // right child
+ }
+ if h.compare(x, h.buf[j]) <= 0 {
+ break
+ }
+ h.buf[i] = h.buf[j]
+ i = j
+ }
+ h.buf[i] = x
+}
diff --git a/heap/heap_test.go b/heap/heap_test.go
new file mode 100644
index 0000000..afbc721
--- /dev/null
+++ b/heap/heap_test.go
@@ -0,0 +1,72 @@
+package heap
+
+import (
+ "cmp"
+ "testing"
+)
+
+func verify(t *testing.T, h *Heap[int], i int) {
+ t.Helper()
+
+ n := h.Len()
+ j1 := 2*i + 1
+ j2 := 2*i + 2
+ if j1 < n {
+ if h.compare(h.buf[j1], h.buf[i]) < 0 {
+ t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h.buf[i], j1, h.buf[j1])
+ return
+ }
+ verify(t, h, j1)
+ }
+ if j2 < n {
+ if h.compare(h.buf[j2], h.buf[i]) < 0 {
+ t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h.buf[i], j1, h.buf[j2])
+ return
+ }
+ verify(t, h, j2)
+ }
+}
+
+func Test(t *testing.T) {
+ t.Parallel()
+
+ h := New(cmp.Compare[int])
+ verify(t, h, 0)
+ h.Grow(20)
+ verify(t, h, 0)
+
+ for i := 20; i > 10; i-- {
+ h.Push(i)
+ }
+ verify(t, h, 0)
+
+ for i := 10; i > 0; i-- {
+ h.Push(i)
+ verify(t, h, 0)
+ }
+
+ for i := 1; h.Len() > 0; i++ {
+ x, ok := h.Pop()
+ if !ok {
+ t.Errorf("Pop() = false, want %d", i)
+ }
+ if i < 20 {
+ h.Push(20 + i)
+ }
+ verify(t, h, 0)
+ if x != i {
+ t.Errorf("%d.th pop got %d; want %d", i, x, i)
+ }
+ }
+}
+
+func TestPopEmpty(t *testing.T) {
+ t.Parallel()
+
+ h := New(cmp.Compare[int])
+ _, gotOk := h.Pop()
+ const want = false
+ if gotOk != want {
+ t.Errorf("Pop() on empty heap = %t, want %t", gotOk, want)
+ }
+}
generated by cgit v1.3.1 (git 2.54.0) at 2026-09-04 05:00:57 +0000