-
Notifications
You must be signed in to change notification settings - Fork 5
/
minheap.go
127 lines (106 loc) · 2.73 KB
/
minheap.go
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
package minheap
import (
"iter"
"github.com/joetifa2003/mm-go/allocator"
"github.com/joetifa2003/mm-go/vector"
)
type MinHeap[T any] struct {
alloc allocator.Allocator
data *vector.Vector[T]
less func(a, b T) bool // check if a < b
}
// New creates a new MinHeap.
func New[T any](alloc allocator.Allocator, less func(a, b T) bool) *MinHeap[T] {
minHeap := allocator.Alloc[MinHeap[T]](alloc)
minHeap.alloc = alloc
minHeap.data = vector.New[T](alloc) // Start with an initial capacity of 16
minHeap.less = less
return minHeap
}
// Push adds a value to the heap.
func (h *MinHeap[T]) Push(value T) {
h.data.Push(value)
h.heapifyUp(h.data.Len() - 1)
}
// Pop removes and returns the minimum value from the heap.
func (h *MinHeap[T]) Pop() T {
if h.data.Len() == 0 {
panic("cannot pop from empty heap")
}
minValue := h.data.RemoveAt(0)
h.heapifyDown(0)
return minValue
}
// Peek returns the minimum value from the heap without removing it.
func (h *MinHeap[T]) Peek() T {
if h.data.Len() == 0 {
panic("cannot peek into empty heap")
}
return h.data.UnsafeAt(0)
}
// Len returns the number of elements in the heap.
func (h *MinHeap[T]) Len() int {
return h.data.Len()
}
// Free frees the heap.
func (h *MinHeap[T]) Free() {
h.data.Free()
allocator.Free(h.alloc, h)
}
// Remove the first element that makes f return true
func (h *MinHeap[T]) Remove(f func(T) bool) {
for i := 0; i < h.data.Len(); i++ {
if f(h.data.UnsafeAt(i)) {
h.removeAt(i)
return
}
}
}
func (h *MinHeap[T]) heapifyUp(index int) {
for index > 0 {
parentIndex := (index - 1) / 2
if h.less(h.data.UnsafeAt(parentIndex), h.data.At(index)) {
break
}
h.swap(parentIndex, index)
index = parentIndex
}
}
func (h *MinHeap[T]) heapifyDown(index int) {
for {
leftChildIndex := 2*index + 1
rightChildIndex := 2*index + 2
smallestIndex := index
if leftChildIndex < h.data.Len() && h.less(h.data.UnsafeAt(leftChildIndex), h.data.At(smallestIndex)) {
smallestIndex = leftChildIndex
}
if rightChildIndex < h.data.Len() && h.less(h.data.UnsafeAt(rightChildIndex), h.data.At(smallestIndex)) {
smallestIndex = rightChildIndex
}
if smallestIndex == index {
break
}
h.swap(index, smallestIndex)
index = smallestIndex
}
}
func (h *MinHeap[T]) swap(i, j int) {
temp := h.data.UnsafeAt(i)
h.data.Set(i, h.data.UnsafeAt(j))
h.data.Set(j, temp)
}
// removeAt removes the element at the specified index from the heap.
func (h *MinHeap[T]) removeAt(index int) {
if index == h.data.Len()-1 {
h.data.Pop()
} else {
h.swap(index, h.data.Len()-1)
h.data.Pop()
h.heapifyDown(index)
h.heapifyUp(index)
}
}
// Iter returns an iterator over the elements of the heap.
func (h *MinHeap[T]) Iter() iter.Seq2[int, T] {
return h.data.Iter()
}