diff options
| author | Rose Hogenson <rosehogenson@posteo.net> | 2025-04-12 17:43:42 -0700 |
|---|---|---|
| committer | Rose Hogenson <rosehogenson@posteo.net> | 2025-04-12 17:43:42 -0700 |
| commit | b3ce67680f630c31bf715447c604019682c9c9ce (patch) | |
| tree | 202b5c98b0b76852f9439785cf665a4b27ad78f0 /heap/heap.go | |
| parent | Tiny performance improvement (diff) | |
| download | deque-b3ce67680f630c31bf715447c604019682c9c9ce.tar.zst | |
Rebrand to github.com/rhogenson/container
I also added a heap package
Diffstat (limited to 'heap/heap.go')
| -rw-r--r-- | heap/heap.go | 87 |
1 files changed, 87 insertions, 0 deletions
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 +} |
