node.gno
12.48 Kb · 488 lines
1package avl
2
3// Node represents a node in an AVL tree.
4type Node struct {
5 key string // key is the unique identifier for the node.
6 value any // value is the data stored in the node.
7 height int8 // height is the height of the node in the tree.
8 size int // size is the number of leaf nodes (key-value pairs) in the subtree rooted at this node.
9 leftNode *Node // leftNode is the left child of the node.
10 rightNode *Node // rightNode is the right child of the node.
11}
12
13// NewNode creates a new node with the given key and value.
14func NewNode(key string, value any) *Node {
15 return &Node{
16 key: key,
17 value: value,
18 height: 0,
19 size: 1,
20 }
21}
22
23// Size returns the size of the subtree rooted at the node.
24func (node *Node) Size() int {
25 if node == nil {
26 return 0
27 }
28 return node.size
29}
30
31// IsLeaf checks if the node is a leaf node (has no children).
32func (node *Node) IsLeaf() bool {
33 return node.height == 0
34}
35
36// Key returns the key of the node.
37func (node *Node) Key() string {
38 return node.key
39}
40
41// Value returns the value of the node.
42func (node *Node) Value() any {
43 return node.value
44}
45
46func (node *Node) _copy() *Node {
47 if node.height == 0 {
48 panic("Why are you copying a value node?")
49 }
50 return &Node{
51 key: node.key,
52 height: node.height,
53 size: node.size,
54 leftNode: node.leftNode,
55 rightNode: node.rightNode,
56 }
57}
58
59// Has checks if a node with the given key exists in the subtree rooted at the node.
60func (node *Node) Has(key string) (has bool) {
61 if node == nil {
62 return false
63 }
64 if node.key == key {
65 return true
66 }
67 if node.height == 0 {
68 return false
69 } else {
70 if key < node.key {
71 return node.getLeftNode().Has(key)
72 } else {
73 return node.getRightNode().Has(key)
74 }
75 }
76}
77
78// Get searches for a node with the given key in the subtree rooted at the node
79// and returns its index, value, and whether it exists.
80func (node *Node) Get(key string) (index int, value any, exists bool) {
81 if node == nil {
82 return 0, nil, false
83 }
84
85 if node.height == 0 {
86 if node.key == key {
87 return 0, node.value, true
88 } else if node.key < key {
89 return 1, nil, false
90 } else {
91 return 0, nil, false
92 }
93 } else {
94 if key < node.key {
95 return node.getLeftNode().Get(key)
96 } else {
97 rightNode := node.getRightNode()
98 index, value, exists = rightNode.Get(key)
99 index += node.size - rightNode.size
100 return index, value, exists
101 }
102 }
103}
104
105// GetByIndex retrieves the key-value pair of the node at the given index
106// in the subtree rooted at the node.
107func (node *Node) GetByIndex(index int) (key string, value any) {
108 if index < 0 {
109 panic("GetByIndex: negative index not allowed")
110 }
111
112 if node.height == 0 {
113 if index != 0 {
114 panic("GetByIndex asked for invalid index")
115 }
116 return node.key, node.value
117 } else {
118 // TODO: could improve this by storing the sizes
119 leftNode := node.getLeftNode()
120 if index < leftNode.size {
121 return leftNode.GetByIndex(index)
122 } else {
123 return node.getRightNode().GetByIndex(index - leftNode.size)
124 }
125 }
126}
127
128// Set inserts a new node with the given key-value pair into the subtree rooted at the node,
129// and returns the new root of the subtree and whether an existing node was updated.
130//
131// XXX consider a better way to do this... perhaps split Node from Node.
132func (node *Node) Set(key string, value any) (newSelf *Node, updated bool) {
133 if node == nil {
134 return NewNode(key, value), false
135 }
136 if node.height == 0 {
137 if key < node.key {
138 return &Node{
139 key: node.key,
140 height: 1,
141 size: 2,
142 leftNode: NewNode(key, value),
143 rightNode: node,
144 }, false
145 } else if key == node.key {
146 return NewNode(key, value), true
147 } else {
148 return &Node{
149 key: key,
150 height: 1,
151 size: 2,
152 leftNode: node,
153 rightNode: NewNode(key, value),
154 }, false
155 }
156 } else {
157 node = node._copy()
158 if key < node.key {
159 node.leftNode, updated = node.getLeftNode().Set(key, value)
160 } else {
161 node.rightNode, updated = node.getRightNode().Set(key, value)
162 }
163 if updated {
164 return node, updated
165 } else {
166 node.calcHeightAndSize()
167 return node.balance(), updated
168 }
169 }
170}
171
172// Remove deletes the node with the given key from the subtree rooted at the node.
173// returns the new root of the subtree, the new leftmost leaf key (if changed),
174// the removed value and the removal was successful.
175func (node *Node) Remove(key string) (
176 newNode *Node, newKey string, value any, removed bool,
177) {
178 if node == nil {
179 return nil, "", nil, false
180 }
181 if node.height == 0 {
182 if key == node.key {
183 return nil, "", node.value, true
184 } else {
185 return node, "", nil, false
186 }
187 } else {
188 if key < node.key {
189 var newLeftNode *Node
190 newLeftNode, newKey, value, removed = node.getLeftNode().Remove(key)
191 if !removed {
192 return node, "", value, false
193 } else if newLeftNode == nil { // left node held value, was removed
194 return node.rightNode, node.key, value, true
195 }
196 node = node._copy()
197 node.leftNode = newLeftNode
198 node.calcHeightAndSize()
199 node = node.balance()
200 return node, newKey, value, true
201 } else {
202 var newRightNode *Node
203 newRightNode, newKey, value, removed = node.getRightNode().Remove(key)
204 if !removed {
205 return node, "", value, false
206 } else if newRightNode == nil { // right node held value, was removed
207 return node.leftNode, "", value, true
208 }
209 node = node._copy()
210 node.rightNode = newRightNode
211 if newKey != "" {
212 node.key = newKey
213 }
214 node.calcHeightAndSize()
215 node = node.balance()
216 return node, "", value, true
217 }
218 }
219}
220
221func (node *Node) getLeftNode() *Node {
222 return node.leftNode
223}
224
225func (node *Node) getRightNode() *Node {
226 return node.rightNode
227}
228
229// rotateRight performs a right rotation on the node and returns the new root.
230// NOTE: overwrites node
231// TODO: optimize balance & rotate
232func (node *Node) rotateRight() *Node {
233 node = node._copy()
234 l := node.getLeftNode()
235 _l := l._copy()
236
237 _lrCached := _l.rightNode
238 _l.rightNode = node
239 node.leftNode = _lrCached
240
241 node.calcHeightAndSize()
242 _l.calcHeightAndSize()
243
244 return _l
245}
246
247// rotateLeft performs a left rotation on the node and returns the new root.
248// NOTE: overwrites node
249// TODO: optimize balance & rotate
250func (node *Node) rotateLeft() *Node {
251 node = node._copy()
252 r := node.getRightNode()
253 _r := r._copy()
254
255 _rlCached := _r.leftNode
256 _r.leftNode = node
257 node.rightNode = _rlCached
258
259 node.calcHeightAndSize()
260 _r.calcHeightAndSize()
261
262 return _r
263}
264
265// calcHeightAndSize updates the height and size of the node based on its children.
266// NOTE: mutates height and size
267func (node *Node) calcHeightAndSize() {
268 node.height = maxInt8(node.getLeftNode().height, node.getRightNode().height) + 1
269 node.size = node.getLeftNode().size + node.getRightNode().size
270}
271
272// calcBalance calculates the balance factor of the node.
273func (node *Node) calcBalance() int {
274 return int(node.getLeftNode().height) - int(node.getRightNode().height)
275}
276
277// balance balances the subtree rooted at the node and returns the new root.
278// NOTE: assumes that node can be modified
279// TODO: optimize balance & rotate
280func (node *Node) balance() (newSelf *Node) {
281 balance := node.calcBalance()
282 if balance > 1 {
283 if node.getLeftNode().calcBalance() >= 0 {
284 // Left Left Case
285 return node.rotateRight()
286 } else {
287 // Left Right Case
288 left := node.getLeftNode()
289 node.leftNode = left.rotateLeft()
290 return node.rotateRight()
291 }
292 }
293 if balance < -1 {
294 if node.getRightNode().calcBalance() <= 0 {
295 // Right Right Case
296 return node.rotateLeft()
297 } else {
298 // Right Left Case
299 right := node.getRightNode()
300 node.rightNode = right.rotateRight()
301 return node.rotateLeft()
302 }
303 }
304 // Nothing changed
305 return node
306}
307
308// Shortcut for TraverseInRange.
309func (node *Node) Iterate(start, end string, cb func(*Node) bool) bool {
310 return node.TraverseInRange(start, end, true, true, cb)
311}
312
313// Shortcut for TraverseInRange.
314func (node *Node) ReverseIterate(start, end string, cb func(*Node) bool) bool {
315 return node.TraverseInRange(start, end, false, true, cb)
316}
317
318// TraverseInRange traverses all nodes, including inner nodes.
319// Start is inclusive and end is exclusive when ascending,
320// Start and end are inclusive when descending.
321// Empty start and empty end denote no start and no end.
322// If leavesOnly is true, only visit leaf nodes.
323// NOTE: To simulate an exclusive reverse traversal,
324// just append 0x00 to start.
325func (node *Node) TraverseInRange(start, end string, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
326 if node == nil {
327 return false
328 }
329 afterStart := (start == "" || start < node.key)
330 startOrAfter := (start == "" || start <= node.key)
331 beforeEnd := false
332 if ascending {
333 beforeEnd = (end == "" || node.key < end)
334 } else {
335 beforeEnd = (end == "" || node.key <= end)
336 }
337
338 // Run callback per inner/leaf node.
339 stop := false
340 if (!node.IsLeaf() && !leavesOnly) ||
341 (node.IsLeaf() && startOrAfter && beforeEnd) {
342 stop = cb(node)
343 if stop {
344 return stop
345 }
346 }
347 if node.IsLeaf() {
348 return stop
349 }
350
351 if ascending {
352 // check lower nodes, then higher
353 if afterStart {
354 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
355 }
356 if stop {
357 return stop
358 }
359 if beforeEnd {
360 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
361 }
362 } else {
363 // check the higher nodes first
364 if beforeEnd {
365 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
366 }
367 if stop {
368 return stop
369 }
370 if afterStart {
371 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
372 }
373 }
374
375 return stop
376}
377
378// TraverseByOffset traverses all nodes, including inner nodes.
379// A limit of math.MaxInt means no limit.
380func (node *Node) TraverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
381 if node == nil {
382 return false
383 }
384
385 // Clamp negative offset to 0; otherwise `delta := first.size - offset`
386 // over-counts and silently drops nodes.
387 if offset < 0 {
388 offset = 0
389 }
390
391 // fast paths. these happen only if TraverseByOffset is called directly on a leaf.
392 if limit <= 0 || offset >= node.size {
393 return false
394 }
395 if node.IsLeaf() {
396 if offset > 0 {
397 return false
398 }
399 return cb(node)
400 }
401
402 // go to the actual recursive function.
403 return node.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
404}
405
406// TraverseByOffset traverses the subtree rooted at the node by offset and limit,
407// in either ascending or descending order, and applies the callback function to each traversed node.
408// If leavesOnly is true, only leaf nodes are visited.
409func (node *Node) traverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
410 // caller guarantees: offset < node.size; limit > 0.
411 if !leavesOnly {
412 if cb(node) {
413 return true // Stop traversal if callback returns true
414 }
415 }
416 first, second := node.getLeftNode(), node.getRightNode()
417 if !ascending {
418 first, second = second, first
419 }
420 if first.IsLeaf() {
421 // either run or skip, based on offset
422 if offset > 0 {
423 offset--
424 } else {
425 if cb(first) {
426 return true // Stop traversal if callback returns true
427 }
428 limit--
429 if limit <= 0 {
430 return true // Stop traversal when limit is reached
431 }
432 }
433 } else {
434 // possible cases:
435 // 1 the offset given skips the first node entirely
436 // 2 the offset skips none or part of the first node, but the limit requires some of the second node.
437 // 3 the offset skips none or part of the first node, and the limit stops our search on the first node.
438 if offset >= first.size {
439 offset -= first.size // 1
440 } else {
441 if first.traverseByOffset(offset, limit, ascending, leavesOnly, cb) {
442 return true
443 }
444 // number of leaves which could actually be called from inside
445 delta := first.size - offset
446 offset = 0
447 if delta >= limit {
448 return true // 3
449 }
450 limit -= delta // 2
451 }
452 }
453
454 // because of the caller guarantees and the way we handle the first node,
455 // at this point we know that limit > 0 and there must be some values in
456 // this second node that we include.
457
458 // => if the second node is a leaf, it has to be included.
459 if second.IsLeaf() {
460 return cb(second)
461 }
462 // => if it is not a leaf, it will still be enough to recursively call this
463 // function with the updated offset and limit
464 return second.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
465}
466
467// Only used in testing...
468func (node *Node) lmd() *Node {
469 if node.height == 0 {
470 return node
471 }
472 return node.getLeftNode().lmd()
473}
474
475// Only used in testing...
476func (node *Node) rmd() *Node {
477 if node.height == 0 {
478 return node
479 }
480 return node.getRightNode().rmd()
481}
482
483func maxInt8(a, b int8) int8 {
484 if a > b {
485 return a
486 }
487 return b
488}
