node.gno
8.71 Kb · 415 lines
1package avl
2
3type Node struct {
4 key string // key is the unique identifier for the node.
5 value any // value is the data stored in the node.
6 height int8 // height is the height of the node in the tree.
7 size int // size is the number of leaf nodes (key-value pairs) in the subtree rooted at this node.
8 leftNode *Node // leftNode is the left child of the node.
9 rightNode *Node // rightNode is the right child of the node.
10}
11
12func NewNode(key string, value any) *Node {
13 return &Node{
14 key: key,
15 value: value,
16 height: 0,
17 size: 1,
18 }
19}
20
21func (node *Node) Size() int {
22 if node == nil {
23 return 0
24 }
25 return node.size
26}
27
28func (node *Node) IsLeaf() bool {
29 return node.height == 0
30}
31
32func (node *Node) Key() string {
33 return node.key
34}
35
36func (node *Node) Value() any {
37 return node.value
38}
39
40func (node *Node) _copy() *Node {
41 if node.height == 0 {
42 panic("Why are you copying a value node?")
43 }
44 return &Node{
45 key: node.key,
46 height: node.height,
47 size: node.size,
48 leftNode: node.leftNode,
49 rightNode: node.rightNode,
50 }
51}
52
53func (node *Node) Has(key string) (has bool) {
54 if node == nil {
55 return false
56 }
57 if node.key == key {
58 return true
59 }
60 if node.height == 0 {
61 return false
62 } else {
63 if key < node.key {
64 return node.getLeftNode().Has(key)
65 } else {
66 return node.getRightNode().Has(key)
67 }
68 }
69}
70
71func (node *Node) Get(key string) (index int, value any, exists bool) {
72 if node == nil {
73 return 0, nil, false
74 }
75
76 if node.height == 0 {
77 if node.key == key {
78 return 0, node.value, true
79 } else if node.key < key {
80 return 1, nil, false
81 } else {
82 return 0, nil, false
83 }
84 } else {
85 if key < node.key {
86 return node.getLeftNode().Get(key)
87 } else {
88 rightNode := node.getRightNode()
89 index, value, exists = rightNode.Get(key)
90 index += node.size - rightNode.size
91 return index, value, exists
92 }
93 }
94}
95
96func (node *Node) GetByIndex(index int) (key string, value any) {
97 if index < 0 {
98 panic("GetByIndex: negative index not allowed")
99 }
100
101 if node.height == 0 {
102 if index != 0 {
103 panic("GetByIndex asked for invalid index")
104 }
105 return node.key, node.value
106 } else {
107 leftNode := node.getLeftNode()
108 if index < leftNode.size {
109 return leftNode.GetByIndex(index)
110 } else {
111 return node.getRightNode().GetByIndex(index - leftNode.size)
112 }
113 }
114}
115
116func (node *Node) Set(key string, value any) (newSelf *Node, updated bool) {
117 if node == nil {
118 return NewNode(key, value), false
119 }
120 if node.height == 0 {
121 if key < node.key {
122 return &Node{
123 key: node.key,
124 height: 1,
125 size: 2,
126 leftNode: NewNode(key, value),
127 rightNode: node,
128 }, false
129 } else if key == node.key {
130 return NewNode(key, value), true
131 } else {
132 return &Node{
133 key: key,
134 height: 1,
135 size: 2,
136 leftNode: node,
137 rightNode: NewNode(key, value),
138 }, false
139 }
140 } else {
141 node = node._copy()
142 if key < node.key {
143 node.leftNode, updated = node.getLeftNode().Set(key, value)
144 } else {
145 node.rightNode, updated = node.getRightNode().Set(key, value)
146 }
147 if updated {
148 return node, updated
149 } else {
150 node.calcHeightAndSize()
151 return node.balance(), updated
152 }
153 }
154}
155
156func (node *Node) Remove(key string) (
157 newNode *Node, newKey string, value any, removed bool,
158) {
159 if node == nil {
160 return nil, "", nil, false
161 }
162 if node.height == 0 {
163 if key == node.key {
164 return nil, "", node.value, true
165 } else {
166 return node, "", nil, false
167 }
168 } else {
169 if key < node.key {
170 var newLeftNode *Node
171 newLeftNode, newKey, value, removed = node.getLeftNode().Remove(key)
172 if !removed {
173 return node, "", value, false
174 } else if newLeftNode == nil { // left node held value, was removed
175 return node.rightNode, node.key, value, true
176 }
177 node = node._copy()
178 node.leftNode = newLeftNode
179 node.calcHeightAndSize()
180 node = node.balance()
181 return node, newKey, value, true
182 } else {
183 var newRightNode *Node
184 newRightNode, newKey, value, removed = node.getRightNode().Remove(key)
185 if !removed {
186 return node, "", value, false
187 } else if newRightNode == nil { // right node held value, was removed
188 return node.leftNode, "", value, true
189 }
190 node = node._copy()
191 node.rightNode = newRightNode
192 if newKey != "" {
193 node.key = newKey
194 }
195 node.calcHeightAndSize()
196 node = node.balance()
197 return node, "", value, true
198 }
199 }
200}
201
202func (node *Node) getLeftNode() *Node {
203 return node.leftNode
204}
205
206func (node *Node) getRightNode() *Node {
207 return node.rightNode
208}
209
210func (node *Node) rotateRight() *Node {
211 node = node._copy()
212 l := node.getLeftNode()
213 _l := l._copy()
214
215 _lrCached := _l.rightNode
216 _l.rightNode = node
217 node.leftNode = _lrCached
218
219 node.calcHeightAndSize()
220 _l.calcHeightAndSize()
221
222 return _l
223}
224
225func (node *Node) rotateLeft() *Node {
226 node = node._copy()
227 r := node.getRightNode()
228 _r := r._copy()
229
230 _rlCached := _r.leftNode
231 _r.leftNode = node
232 node.rightNode = _rlCached
233
234 node.calcHeightAndSize()
235 _r.calcHeightAndSize()
236
237 return _r
238}
239
240func (node *Node) calcHeightAndSize() {
241 node.height = maxInt8(node.getLeftNode().height, node.getRightNode().height) + 1
242 node.size = node.getLeftNode().size + node.getRightNode().size
243}
244
245func (node *Node) calcBalance() int {
246 return int(node.getLeftNode().height) - int(node.getRightNode().height)
247}
248
249func (node *Node) balance() (newSelf *Node) {
250 balance := node.calcBalance()
251 if balance > 1 {
252 if node.getLeftNode().calcBalance() >= 0 {
253 return node.rotateRight()
254 } else {
255 left := node.getLeftNode()
256 node.leftNode = left.rotateLeft()
257 return node.rotateRight()
258 }
259 }
260 if balance < -1 {
261 if node.getRightNode().calcBalance() <= 0 {
262 return node.rotateLeft()
263 } else {
264 right := node.getRightNode()
265 node.rightNode = right.rotateRight()
266 return node.rotateLeft()
267 }
268 }
269 return node
270}
271
272func (node *Node) Iterate(start, end string, cb func(*Node) bool) bool {
273 return node.TraverseInRange(start, end, true, true, cb)
274}
275
276func (node *Node) ReverseIterate(start, end string, cb func(*Node) bool) bool {
277 return node.TraverseInRange(start, end, false, true, cb)
278}
279
280func (node *Node) TraverseInRange(start, end string, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
281 if node == nil {
282 return false
283 }
284 afterStart := (start == "" || start < node.key)
285 startOrAfter := (start == "" || start <= node.key)
286 beforeEnd := false
287 if ascending {
288 beforeEnd = (end == "" || node.key < end)
289 } else {
290 beforeEnd = (end == "" || node.key <= end)
291 }
292
293 stop := false
294 if (!node.IsLeaf() && !leavesOnly) ||
295 (node.IsLeaf() && startOrAfter && beforeEnd) {
296 stop = cb(node)
297 if stop {
298 return stop
299 }
300 }
301 if node.IsLeaf() {
302 return stop
303 }
304
305 if ascending {
306 if afterStart {
307 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
308 }
309 if stop {
310 return stop
311 }
312 if beforeEnd {
313 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
314 }
315 } else {
316 if beforeEnd {
317 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
318 }
319 if stop {
320 return stop
321 }
322 if afterStart {
323 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
324 }
325 }
326
327 return stop
328}
329
330func (node *Node) TraverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
331 if node == nil {
332 return false
333 }
334
335 if offset < 0 {
336 offset = 0
337 }
338
339 if limit <= 0 || offset >= node.size {
340 return false
341 }
342 if node.IsLeaf() {
343 if offset > 0 {
344 return false
345 }
346 return cb(node)
347 }
348
349 return node.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
350}
351
352func (node *Node) traverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
353 if !leavesOnly {
354 if cb(node) {
355 return true // Stop traversal if callback returns true
356 }
357 }
358 first, second := node.getLeftNode(), node.getRightNode()
359 if !ascending {
360 first, second = second, first
361 }
362 if first.IsLeaf() {
363 if offset > 0 {
364 offset--
365 } else {
366 if cb(first) {
367 return true // Stop traversal if callback returns true
368 }
369 limit--
370 if limit <= 0 {
371 return true // Stop traversal when limit is reached
372 }
373 }
374 } else {
375 if offset >= first.size {
376 offset -= first.size // 1
377 } else {
378 if first.traverseByOffset(offset, limit, ascending, leavesOnly, cb) {
379 return true
380 }
381 delta := first.size - offset
382 offset = 0
383 if delta >= limit {
384 return true // 3
385 }
386 limit -= delta // 2
387 }
388 }
389
390 if second.IsLeaf() {
391 return cb(second)
392 }
393 return second.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
394}
395
396func (node *Node) lmd() *Node {
397 if node.height == 0 {
398 return node
399 }
400 return node.getLeftNode().lmd()
401}
402
403func (node *Node) rmd() *Node {
404 if node.height == 0 {
405 return node
406 }
407 return node.getRightNode().rmd()
408}
409
410func maxInt8(a, b int8) int8 {
411 if a > b {
412 return a
413 }
414 return b
415}
