-
Notifications
You must be signed in to change notification settings - Fork 0
/
node_heap.go
163 lines (152 loc) · 3.76 KB
/
node_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
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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
package main
import (
"bytes"
"errors"
"fmt"
"github.com/disiqueira/gotree"
)
// balanced binary heap of Nodes sorted by K
type NodeHeap struct {
Arr []*Node
}
func NewNodeHeap() *NodeHeap {
// 0 is a nil element not considered, because it makes the
// array shifting math cleaner
return &NodeHeap{Arr: []*Node{nil}}
}
func (h *NodeHeap) Add(n *Node) (ix int) {
// append to array end (will bubble up below)
h.Arr = append(h.Arr, n)
ix = len(h.Arr) - 1
// compute F = G + H
n.F = n.G + n.H
// bubble up (needs HeapIX set)
n.HeapIX = ix
ix = h.bubbleUp(ix)
// return ix to user
return ix
}
func (h *NodeHeap) bubbleUp(ix int) int {
// if ix > 1 and K(ix) < K(ix>>1)
// (if we're not the top node and our K value is less than the parent)
for ix > 1 && h.Arr[ix].F < h.Arr[ix>>1].F {
// swap Nodes in heap
h.Arr[ix], h.Arr[ix>>1] = h.Arr[ix>>1], h.Arr[ix]
// swap HeapIX for nodes
h.Arr[ix].HeapIX, h.Arr[ix>>1].HeapIX =
h.Arr[ix>>1].HeapIX, h.Arr[ix].HeapIX
// set ix for next iter to be the ix we just swapped to
ix = ix >> 1
}
return ix
}
func (h *NodeHeap) bubbleDown(ix int) int {
for {
// beginning state: lesser node is the one we're on
lesser := ix
// l index, r index
lix := (ix << 1)
rix := (ix << 1) + 1
// if we've reached the bottom, return
if !(lix < len(h.Arr) || rix < len(h.Arr)) {
return ix
}
// if left child exists and is lesser, set lesser to lix
if lix < len(h.Arr) && h.Arr[lix].F < h.Arr[lesser].F {
lesser = lix
}
// if left child exists and is lesser, set lesser to rix
if rix < len(h.Arr) && h.Arr[rix].F < h.Arr[lesser].F {
lesser = rix
}
// if one of children was lesser, swap and continue bubble down
// from that node
if lesser != ix {
// swap Nodes in heap
h.Arr[ix], h.Arr[lesser] = h.Arr[lesser], h.Arr[ix]
// swap HeapIX for nodes
h.Arr[ix].HeapIX, h.Arr[lesser].HeapIX =
h.Arr[lesser].HeapIX, h.Arr[ix].HeapIX
// continue to bubble down
ix = lesser
continue
} else {
// else, no child was lesser. we can stop here
return ix
}
}
}
func (h *NodeHeap) Pop() (*Node, error) {
if h.Len() == 0 {
return nil, errors.New("heap empty")
}
// get root elem and replace with last element
n := h.Arr[1]
last_ix := len(h.Arr) - 1
// bring last element to root, bubble new root node down if
// needed, before return
h.Arr[1] = h.Arr[last_ix]
h.Arr[1].HeapIX = 1
h.Arr = h.Arr[:last_ix]
h.bubbleDown(1)
return n, nil
}
func (h *NodeHeap) Modified(n *Node) {
// if less than parent, bubble up
parentIX := n.HeapIX >> 1
if parentIX > 0 {
parent := h.Arr[parentIX]
if n.F < parent.F {
h.bubbleUp(n.HeapIX)
return
}
}
// if greater than either child, bubble down
lix := (n.HeapIX << 1)
rix := (n.HeapIX << 1) + 1
for _, ix := range []int{lix, rix} {
if ix < len(h.Arr) {
child := h.Arr[ix]
if n.F > child.F {
h.bubbleDown(n.HeapIX)
return
}
}
}
}
func (h *NodeHeap) Len() int {
return len(h.Arr) - 1
}
func (h *NodeHeap) Clear() {
h.Arr = h.Arr[:1]
}
func (h *NodeHeap) String() string {
// for building string
var buffer bytes.Buffer
// print array
// if elements, print tree
if len(h.Arr) > 1 {
buffer.WriteString("\n")
// build tree using gotree package by descending recursively
var addChildren func(node gotree.Tree, ix int)
addChildren = func(node gotree.Tree, ix int) {
rix := (ix << 1) + 1
if rix < len(h.Arr) {
r := node.Add(fmt.Sprintf("[%d]%s",
rix, h.Arr[rix].String()))
addChildren(r, rix)
}
lix := (ix << 1)
if lix < len(h.Arr) {
l := node.Add(fmt.Sprintf("[%d]%s",
lix, h.Arr[lix].String()))
addChildren(l, lix)
}
}
tree := gotree.New(fmt.Sprintf("[%d]%s",
1, h.Arr[1].String()))
addChildren(tree, 1)
buffer.WriteString(tree.Print())
}
return buffer.String()
}