blob: 612d5a8445de15ccff56b0b8d59eba7c261c9209 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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
}
|
