aboutsummaryrefslogtreecommitdiff
path: root/internal/heap/heap.go
diff options
context:
space:
mode:
Diffstat (limited to 'internal/heap/heap.go')
-rw-r--r--internal/heap/heap.go93
1 files changed, 93 insertions, 0 deletions
diff --git a/internal/heap/heap.go b/internal/heap/heap.go
new file mode 100644
index 00000000..7ad551ab
--- /dev/null
+++ b/internal/heap/heap.go
@@ -0,0 +1,93 @@
+package heap
+
+// Heap is one slice-backed binary heap.
+type Heap[T any] struct {
+ items []T
+ less func(left, right T) bool
+}
+
+// New builds one empty heap ordered by less.
+func New[T any](less func(left, right T) bool) *Heap[T] {
+ return &Heap[T]{less: less}
+}
+
+// Len reports the number of queued items.
+func (heap *Heap[T]) Len() int {
+ return len(heap.items)
+}
+
+// Empty reports whether the heap is empty.
+func (heap *Heap[T]) Empty() bool {
+ return len(heap.items) == 0
+}
+
+// Peek returns the top item without removing it.
+func (heap *Heap[T]) Peek() (T, bool) {
+ if len(heap.items) == 0 {
+ var zero T
+
+ return zero, false
+ }
+
+ return heap.items[0], true
+}
+
+// Push inserts one item.
+func (heap *Heap[T]) Push(item T) {
+ heap.items = append(heap.items, item)
+ heap.siftUp(len(heap.items) - 1)
+}
+
+// Pop removes one top item.
+func (heap *Heap[T]) Pop() (T, bool) {
+ if len(heap.items) == 0 {
+ var zero T
+
+ return zero, false
+ }
+
+ last := len(heap.items) - 1
+ top := heap.items[0]
+ heap.items[0] = heap.items[last]
+ heap.items = heap.items[:last]
+
+ if len(heap.items) > 0 {
+ heap.siftDown(0)
+ }
+
+ return top, true
+}
+
+func (heap *Heap[T]) siftUp(idx int) {
+ for idx > 0 {
+ parent := (idx - 1) / 2
+ if !heap.less(heap.items[idx], heap.items[parent]) {
+ return
+ }
+
+ heap.items[idx], heap.items[parent] = heap.items[parent], heap.items[idx]
+ idx = parent
+ }
+}
+
+func (heap *Heap[T]) siftDown(idx int) {
+ for {
+ left := idx*2 + 1
+ if left >= len(heap.items) {
+ return
+ }
+
+ best := left
+ right := left + 1
+ if right < len(heap.items) && heap.less(heap.items[right], heap.items[left]) {
+ best = right
+ }
+
+ if !heap.less(heap.items[best], heap.items[idx]) {
+ return
+ }
+
+ heap.items[idx], heap.items[best] = heap.items[best], heap.items[idx]
+ idx = best
+ }
+}