node.gno
12.33 Kb · 482 lines
1package avl
2
3//----------------------------------------
4// Node
5
6// Node represents a node in an AVL tree.
7type Node struct {
8 key string // key is the unique identifier for the node.
9 value any // value is the data stored in the node.
10 height int8 // height is the height of the node in the tree.
11 size int // size is the number of leaf nodes (key-value pairs) in the subtree rooted at this node.
12 leftNode *Node // leftNode is the left child of the node.
13 rightNode *Node // rightNode is the right child of the node.
14}
15
16// NewNode creates a new node with the given key and value.
17func NewNode(key string, value any) *Node {
18 return &Node{
19 key: key,
20 value: value,
21 height: 0,
22 size: 1,
23 }
24}
25
26// Size returns the size of the subtree rooted at the node.
27func (node *Node) Size() int {
28 if node == nil {
29 return 0
30 }
31 return node.size
32}
33
34// IsLeaf checks if the node is a leaf node (has no children).
35func (node *Node) IsLeaf() bool {
36 return node.height == 0
37}
38
39// Key returns the key of the node.
40func (node *Node) Key() string {
41 return node.key
42}
43
44// Value returns the value of the node.
45func (node *Node) Value() any {
46 return node.value
47}
48
49func (node *Node) _copy() *Node {
50 if node.height == 0 {
51 panic("Why are you copying a value node?")
52 }
53 return &Node{
54 key: node.key,
55 height: node.height,
56 size: node.size,
57 leftNode: node.leftNode,
58 rightNode: node.rightNode,
59 }
60}
61
62// Has checks if a node with the given key exists in the subtree rooted at the node.
63func (node *Node) Has(key string) (has bool) {
64 if node == nil {
65 return false
66 }
67 if node.key == key {
68 return true
69 }
70 if node.height == 0 {
71 return false
72 } else {
73 if key < node.key {
74 return node.getLeftNode().Has(key)
75 } else {
76 return node.getRightNode().Has(key)
77 }
78 }
79}
80
81// Get searches for a node with the given key in the subtree rooted at the node
82// and returns its index, value, and whether it exists.
83func (node *Node) Get(key string) (index int, value any, exists bool) {
84 if node == nil {
85 return 0, nil, false
86 }
87
88 if node.height == 0 {
89 if node.key == key {
90 return 0, node.value, true
91 } else if node.key < key {
92 return 1, nil, false
93 } else {
94 return 0, nil, false
95 }
96 } else {
97 if key < node.key {
98 return node.getLeftNode().Get(key)
99 } else {
100 rightNode := node.getRightNode()
101 index, value, exists = rightNode.Get(key)
102 index += node.size - rightNode.size
103 return index, value, exists
104 }
105 }
106}
107
108// GetByIndex retrieves the key-value pair of the node at the given index
109// in the subtree rooted at the node.
110func (node *Node) GetByIndex(index int) (key string, value any) {
111 if node.height == 0 {
112 if index == 0 {
113 return node.key, node.value
114 } else {
115 panic("GetByIndex asked for invalid index")
116 }
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 // fast paths. these happen only if TraverseByOffset is called directly on a leaf.
386 if limit <= 0 || offset >= node.size {
387 return false
388 }
389 if node.IsLeaf() {
390 if offset > 0 {
391 return false
392 }
393 return cb(node)
394 }
395
396 // go to the actual recursive function.
397 return node.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
398}
399
400// TraverseByOffset traverses the subtree rooted at the node by offset and limit,
401// in either ascending or descending order, and applies the callback function to each traversed node.
402// If leavesOnly is true, only leaf nodes are visited.
403func (node *Node) traverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
404 // caller guarantees: offset < node.size; limit > 0.
405 if !leavesOnly {
406 if cb(node) {
407 return true // Stop traversal if callback returns true
408 }
409 }
410 first, second := node.getLeftNode(), node.getRightNode()
411 if !ascending {
412 first, second = second, first
413 }
414 if first.IsLeaf() {
415 // either run or skip, based on offset
416 if offset > 0 {
417 offset--
418 } else {
419 if cb(first) {
420 return true // Stop traversal if callback returns true
421 }
422 limit--
423 if limit <= 0 {
424 return true // Stop traversal when limit is reached
425 }
426 }
427 } else {
428 // possible cases:
429 // 1 the offset given skips the first node entirely
430 // 2 the offset skips none or part of the first node, but the limit requires some of the second node.
431 // 3 the offset skips none or part of the first node, and the limit stops our search on the first node.
432 if offset >= first.size {
433 offset -= first.size // 1
434 } else {
435 if first.traverseByOffset(offset, limit, ascending, leavesOnly, cb) {
436 return true
437 }
438 // number of leaves which could actually be called from inside
439 delta := first.size - offset
440 offset = 0
441 if delta >= limit {
442 return true // 3
443 }
444 limit -= delta // 2
445 }
446 }
447
448 // because of the caller guarantees and the way we handle the first node,
449 // at this point we know that limit > 0 and there must be some values in
450 // this second node that we include.
451
452 // => if the second node is a leaf, it has to be included.
453 if second.IsLeaf() {
454 return cb(second)
455 }
456 // => if it is not a leaf, it will still be enough to recursively call this
457 // function with the updated offset and limit
458 return second.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
459}
460
461// Only used in testing...
462func (node *Node) lmd() *Node {
463 if node.height == 0 {
464 return node
465 }
466 return node.getLeftNode().lmd()
467}
468
469// Only used in testing...
470func (node *Node) rmd() *Node {
471 if node.height == 0 {
472 return node
473 }
474 return node.getRightNode().rmd()
475}
476
477func maxInt8(a, b int8) int8 {
478 if a > b {
479 return a
480 }
481 return b
482}
