-
Notifications
You must be signed in to change notification settings - Fork 0
/
heap.go
94 lines (82 loc) · 1.45 KB
/
heap.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
package heap
type Interface[T any] interface {
~[]T
Less(i, j int) bool
}
type Heap[T any, V Interface[T]] []T
func NewHeap[T any, V Interface[T]](arr V) Heap[T, V] {
heap := Heap[T, V](arr)
for i := 1; i < len(arr); i++ {
heap.up(i)
}
return heap
}
func (heap Heap[T, V]) Init() {
// heapify
n := len(heap)
for i := n/2 - 1; i >= 0; i-- {
heap.down(i, n)
}
}
func (heap *Heap[T, V]) Push(x T) {
h := *heap
*heap = append(h, x)
heap.up(len(h))
}
func (heap *Heap[T, V]) Pop() T {
h := *heap
n := len(h) - 1
item := h[0]
h[0], h[n] = h[n], h[0]
heap.down(0, n)
*heap = h[:n]
return item
}
func (heap Heap[T, V]) First() T {
return heap[0]
}
func (heap Heap[T, V]) Last() T {
return heap[len(heap)-1]
}
func (heap *Heap[T, V]) Remove(i int) T {
h := *heap
n := len(h) - 1
item := h[i]
if n != i {
h[i], h[n] = h[n], h[i]
if !heap.down(i, n) {
heap.up(i)
}
}
*heap = h[:n]
return item
}
func (heap Heap[T, V]) up(j int) {
for {
i := (j - 1) / 2 // parent
if i == j || !V(heap).Less(j, i) {
break
}
heap[i], heap[j] = heap[j], heap[i]
j = i
}
}
func (heap Heap[T, V]) down(i0, n int) bool {
i := i0
for {
j1 := 2*i + 1
if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
break
}
j := j1 // left child
if j2 := j1 + 1; j2 < n && V(heap).Less(j2, j1) {
j = j2 // = 2*i + 2 // right child
}
if !V(heap).Less(j, i) {
break
}
heap[i], heap[j] = heap[j], heap[i]
i = j
}
return i > i0
}